def add_graph_struct(self, graph_name: str, graph_struct: GraphStruct) -> int:
graph_view = self.graphs.collection_view(graph_name)
graph_struct.graph_name = graph_name
try:
new_revision = self.read_graph_struct(graph_name).revision + 1
except:
new_revision = 0
graph_struct.revision = new_revision
# FIXME(luckygeck): possible race condition
graph_view.put(str(new_revision), graph_struct.to_json())
return new_revision
def read_graph_struct(self, graph_name: str, revision: int = -1) -> GraphStruct:
graph_view = self.graphs.collection_view(graph_name)
if revision == -1:
last_revision_struct = max(
map(lambda _: GraphStruct.create(_[1]), graph_view.iterate_all(include_value=True)),
key=lambda x: x.revision,
default=None
)
if last_revision_struct is None:
raise GraphStructureNotFound(graph_name)
return last_revision_struct
return GraphStruct.create(graph_view.get(str(revision)))
def _prepare_graph_struct(name: Optional[str], graph: TextIO, hosts: List[str], graph_format: str) -> dict:
if graph_format == 'raw':
return json.load(graph)
assert name and hosts, 'Only raw graph format can not set hosts and name'
result = GraphStruct()
result.graph_name = name
result.clusters.from_json({'I': hosts})
if graph_format == 'script':
task = ExtendedTaskStruct()
task.task_name = 'main'
task.hosts.append('I')
task.task_struct.executor.name = 'shell'
executor_cfg = ShellExecutorConfig()
executor_cfg.shell_script = graph.read()
task.task_struct.executor.config = executor_cfg.to_json()
result.tasks.from_json([task.to_json()])
elif graph_format == 'makefile':
raise NotImplementedError()
return result.to_json()
def list_graphs(self, graph_name: Optional[str] = None, offset: int = 0, limit: Optional[int] = None,
with_info: bool = False) -> List[GraphStruct]:
"""List all known graphs. If graph_name is set, then include only all versions of this graph.
If with_info is false, only graph_name and revision are received.
:returns List[GraphStruct]: List of received graph structs
"""
data = {offset: offset, with_info: '1' if with_info else '0'}
if graph_name:
data['graph_name'] = graph_name
if limit is not None:
data['limit'] = limit
return [GraphStruct.create(_) for _ in requests.get(self._url_prefix + 'graphs', params=data).json()['payload']]
def add_graph_struct(self, graph_name: str, graph_struct: dict) -> int:
return self.backend.add_graph_struct(graph_name, GraphStruct.create(graph_struct))
def list_graph_struct(self, graph_name: Optional[str] = None, with_info: bool = False) -> Iterator[
Tuple[str, int, Optional[GraphStruct]]]:
db = self.graphs.collection_view(graph_name) if graph_name else self.graphs
for key, graph_struct in db.iterate_all(include_value=with_info):
name, revision = (graph_name, key) if graph_name else key.split('=', 1)
yield name, revision, GraphStruct.create(graph_struct) if graph_struct else None
def read_graph(self, graph_name: str, graph_revision: Optional[int] = None) -> GraphStruct:
"""Read info about a specified graph's version. If graph_revision is not set, last revision is used.
:returns GraphStruct: info about a graph
"""
return GraphStruct.create(requests.get(self._get_graph_url(graph_name, graph_revision)).json()['payload'])
