feat: add KG downloading (#71)

Signed-off-by: Yusik Kim <107410898+kmyusk@users.noreply.github.com>
DS4SD · Mar 14, 2023 · 77cccc0 · 77cccc0
1 parent 3e91f63
commit 77cccc0
Show file tree

Hide file tree

Showing 6 changed files with 263 additions and 1 deletion.
diff --git a/deepsearch/cps/apis/public/api/knowledge_graphs_api.py b/deepsearch/cps/apis/public/api/knowledge_graphs_api.py
diff --git a/deepsearch/cps/cli/kgs.py b/deepsearch/cps/cli/kgs.py
@@ -109,5 +109,24 @@ def save_snapshot_of_data_flow(
         typer.echo(f"A snapshot is being taken, see task {task.task_id!r}")
 
 
+@app.command(name="download", help="Download KG")
+def download_knowledge_graph(
+    proj_key: str = PROJ_KEY,
+    kg_key: str = KG_KEY,
+    output: OutputEnum = OutputOption,
+):
+    api = CpsApi.default_from_env()
+    kg = api.knowledge_graphs.get(project=proj_key, key=kg_key)
+    if kg is None:
+        raise typer.BadParameter(
+            f"Unknown Knowledge Graph {kg_key!r} in project {proj_key!r}"
+        )
+    url = kg.download()
+
+    results = [{"url": url}]
+    # TODO: augment with topology details
+    cli_output(results, output, headers="keys")
+
+
 if __name__ == "__main__":
     app()
diff --git a/deepsearch/cps/client/components/knowledge_graphs.py b/deepsearch/cps/client/components/knowledge_graphs.py
@@ -136,3 +136,12 @@ def save_snapshot_of_data_flow(
         )
 
         return task
+
+    def download(self):
+        task: ProjectTask = (
+            self.api.knowledge_graphs.sw_api.download_project_knowledge_graph(
+                self.project, self.key
+            )
+        )
+
+        return self.api.tasks.wait_for(self.project, task.task_id)
diff --git a/docs/apis/public/KnowledgeGraphsApi.md b/docs/apis/public/KnowledgeGraphsApi.md
diff --git a/docs/guide/kgs.md b/docs/guide/kgs.md
@@ -52,3 +52,32 @@ result = api.queries.run(query)
 ```
 
 More example queries are available in the [examples gallery](../gallery/index.md).
+
+### Download a KG
+
+=== "CLI"
+
+    Using the [`deepsearch cps`](../cli-reference.md#cps) component:
+    <div class="termy">
+
+    ```console
+    $ deepsearch cps kgs download -p 'abc123' -k 'zxc987'
+
+    url
+    -------------------------------- 
+    https://host.domain/download-url
+    ```
+
+    </div>
+
+=== "Python"
+
+    After you have generated the `api` object (from [login configuration](../getting_started/#authentication))
+
+    ```python
+    proj_key = "abc123" # get your proj_key from the UI details
+    kg_key = "xyz987" # get kg_key (also called bag_key) from the UI details
+    kg = api.knowledge_graphs.get(proj_key, kg_key)
+
+    download_url = kg.download()
+    ```
diff --git a/tools/swagger-client-generator/swagger-cps.json b/tools/swagger-client-generator/swagger-cps.json