From ff617d1f8dbbf75063863079bc5c5a8ce86517c4 Mon Sep 17 00:00:00 2001
From: Doug Blank <doug.blank@gmail.com>
Date: Sun, 14 Jan 2024 09:04:57 -0500
Subject: [PATCH] Changed remote strategy; allow .db datagrid names; removed
 --root

---
 backend/kangas/cli/server.py         | 13 +++----------
 backend/kangas/integrations/comet.py |  6 +++++-
 backend/kangas/server/queries.py     | 29 +++++++++++++++++-----------
 3 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/backend/kangas/cli/server.py b/backend/kangas/cli/server.py
index fe822ce..cb7b837 100644
--- a/backend/kangas/cli/server.py
+++ b/backend/kangas/cli/server.py
@@ -61,13 +61,6 @@ def get_parser_arguments(parser):
         type=int,
         default=None,
     )
-    parser.add_argument(
-        "-r",
-        "--root",
-        help="The directory from which to server datagrid files; also can use KANGAS_ROOT env variable",
-        type=str,
-        default=None,
-    )
     parser.add_argument(
         "-b",
         "--backend",
@@ -227,9 +220,8 @@ def server(parsed_args, remaining=None):
         terminate()
         return
 
-    print(
-        "Serving DataGrids from directory: %r"
-        % (parsed_args.root or kangas.server.KANGAS_ROOT)
+    print("Serving DataGrids from directory: %r"
+        % (kangas.server.KANGAS_ROOT)
     )
 
     if parsed_args.frontend != "no":
@@ -257,6 +249,7 @@ def server(parsed_args, remaining=None):
                 "KANGAS_BACKEND_HOST": str(KANGAS_BACKEND_HOST),
                 "KANGAS_BACKEND_PROTOCOL": KANGAS_BACKEND_PROTOCOL,
                 "KANGAS_HIDE_SELECTOR": str(KANGAS_HIDE_SELECTOR),
+                "KANGAS_ROOT": kangas.server.KANGAS_ROOT
             }
         )
         # Only add these if they are set:
diff --git a/backend/kangas/integrations/comet.py b/backend/kangas/integrations/comet.py
index a2a0daf..9acab73 100644
--- a/backend/kangas/integrations/comet.py
+++ b/backend/kangas/integrations/comet.py
@@ -229,7 +229,11 @@ def export_to_comet(path, name, options):
     ):
         asset_metadata = json.loads(asset_metadata_string)
         asset_metadata["source"] = asset_map[asset_id]["web"]
-        asset_metadata["cometAssetId"] = asset_map[asset_id]["assetId"]
+        asset_metadata["remote"] = {
+            "framework": "comet",
+            "assetId": asset_map[asset_id]["assetId"],
+            "experimentId": experiment.id,
+        }
         cur.execute(
             "UPDATE assets SET asset_metadata = ? WHERE asset_id = ?;",
             (json.dumps(asset_metadata), asset_id),
diff --git a/backend/kangas/server/queries.py b/backend/kangas/server/queries.py
index 1c3de59..5d6a9f1 100644
--- a/backend/kangas/server/queries.py
+++ b/backend/kangas/server/queries.py
@@ -2447,7 +2447,7 @@ def select_asset(dgid, asset_id, thumbnail=False, return_image=False):
     cur = conn.cursor()
     selection = (
         "SELECT asset_data, asset_type, asset_thumbnail, "
-        + 'json_extract(asset_metadata, "$.source") as asset_source, '
+        + 'json_extract(asset_metadata, "$.remote") as asset_remote, '
         + 'json_extract(asset_metadata, "$.annotations") as asset_annotations '
         + 'from assets where asset_id = "{asset_id}";'
     )
@@ -2459,17 +2459,24 @@ def select_asset(dgid, asset_id, thumbnail=False, return_image=False):
     LOGGER.debug("SQL %s seconds", time.time() - start_time)
 
     if row:
-        asset_data, asset_type, asset_thumbnail, asset_source, asset_annotations = row
-        if asset_source:  # FIXME: asset_type == ["Image"]
+        asset_data, asset_type, asset_thumbnail, asset_remote, asset_annotations = row
+        if asset_remote:  
+            # FIXME: asset_type == ["Image"]
             # FIXME: move to Image class
             # FIXME: use a cache?
-            url_data = urllib.request.urlopen(asset_source)
-            with io.BytesIO() as fp:
-                fp.write(url_data.read())
-                image = PIL.Image.open(fp)
-                if image.mode == "CMYK":
-                    image = image.convert("RGB")
-                asset_data = image_to_fp(image, "png").read()
+            remote = json.loads(asset_remote)
+            experiment_key = remote["experimentId"]
+            asset_id = remote["assetId"]
+            if remote["framework"] == "comet":
+                import comet_ml
+                api = comet_ml.API()
+                asset_data = api._client.get_experiment_asset(
+                    asset_id=asset_id,
+                    experiment_key=experiment_key,
+                    return_type="binary",
+                )
+            else:
+                raise Exception("Unknown remote type")
 
         if thumbnail and asset_type in ["Image"]:
             if asset_annotations:
@@ -2534,7 +2541,7 @@ def list_datagrids():
         filename
         for directory, dirs, files in walk(KANGAS_ROOT)
         for filename in files
-        if filename.endswith(".datagrid")
+        if filename.endswith(".datagrid") or filename.endswith(".db")
     ]
 
     return [