blob: 7acce4120b44490c10efe424de41d50d7940b81e [file] [log] [blame]
Scott Bakerbba67b62019-01-28 17:38:21 -08001# Copyright 2017-present Open Networking Foundation
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import os
16from xosconfig import Config
17from multistructlog import create_logger
18
19log = create_logger(Config().get("logging"))
20
21
22class ModelLoadClient(object):
Scott Baker7ff8ad92019-02-15 17:02:41 -080023 REQUIRE_CLEAN = 0
24 AUTOMATICALLY_CLEAN = 1
25 PURGE = 2
26
Scott Bakerbba67b62019-01-28 17:38:21 -080027 def __init__(self, api):
28 self.api = api
29
30 def upload_models(self, name, dir, version="unknown"):
31 request = self.api.dynamicload_pb2.LoadModelsRequest(name=name, version=version)
32
33 for fn in os.listdir(dir):
34 if fn.endswith(".xproto"):
35 item = request.xprotos.add()
36 item.filename = fn
37 item.contents = open(os.path.join(dir, fn)).read()
38
39 models_fn = os.path.join(dir, "models.py")
40 if os.path.exists(models_fn):
41 item = request.decls.add()
42 item.filename = "models.py"
43 item.contents = open(models_fn).read()
44
45 attic_dir = os.path.join(dir, "attic")
46 if os.path.exists(attic_dir):
47 log.warn(
48 "Attics are deprecated, please use the legacy=True option in xProto"
49 )
50 for fn in os.listdir(attic_dir):
51 if fn.endswith(".py"):
52 item = request.attics.add()
53 item.filename = fn
54 item.contents = open(os.path.join(attic_dir, fn)).read()
55
56 api_convenience_dir = os.path.join(dir, "convenience")
57 if os.path.exists(api_convenience_dir):
58 for fn in os.listdir(api_convenience_dir):
59 if fn.endswith(".py") and "test" not in fn:
60 item = request.convenience_methods.add()
61 item.filename = fn
62 item.contents = open(os.path.join(api_convenience_dir, fn)).read()
63
Scott Baker8cc464d2019-02-06 16:51:42 -080064 # migrations directory is a sibling to the models directory
65 migrations_dir = os.path.join(dir, "..", "migrations")
66 if os.path.exists(migrations_dir):
67 for fn in os.listdir(migrations_dir):
68 if fn.endswith(".py") and "test" not in fn:
69 item = request.migrations.add()
70 item.filename = fn
71 item.contents = open(os.path.join(migrations_dir, fn)).read()
72
Scott Bakerbba67b62019-01-28 17:38:21 -080073 result = self.api.dynamicload.LoadModels(request)
Scott Baker7ff8ad92019-02-15 17:02:41 -080074
75 return result
76
77 def unload_models(self, name, version="unknown", cleanup_behavior=REQUIRE_CLEAN):
78 request = self.api.dynamicload_pb2.UnloadModelsRequest(
79 name=name,
80 version=version,
81 cleanup_behavior=cleanup_behavior)
82 result = self.api.dynamicload.UnloadModels(request)
83
84 return result
85