Merge pull request #11 from compute-tooling/private

Make cloud storage files private

cs_storage/__init__.py

@@ -225,16 +225,12 @@ def write(task_id, loc_result, do_upload=True, protocol="gcs"):
 
 
 def read(rem_result, json_serializable=True, protocol="gcs"):
-    # compute studio results have public read access.
-    # fs = gcsfs.GCSFileSystem(token="anon")
     s = time.time()
     RemoteResult().load(rem_result)
     read = {"renderable": [], "downloadable": []}
     for category in rem_result:
         with fs.open(
-            f"{protocol}://{BUCKET}/{rem_result[category]['ziplocation']}",
-            "rb",
-            **{protocol: {"token": "anon"}},
+            f"{protocol}://{BUCKET}/{rem_result[category]['ziplocation']}", "rb",
         ) as f:
             res = f.read()
 
@@ -259,9 +255,14 @@ def read(rem_result, json_serializable=True, protocol="gcs"):
     return read
 
 
+def read_screenshot(screenshot_id, protocol="gcs"):
+    if not screenshot_id.endswith(".png"):
+        screenshot_id += ".png"
+    with fs.open(f"{protocol}://{BUCKET}/{screenshot_id}", "rb") as f:
+        return f.read()
+
+
 def add_screenshot_links(rem_result):
     for rem_output in rem_result["renderable"]["outputs"]:
-        rem_output[
-            "screenshot"
-        ] = f"https://storage.googleapis.com/{BUCKET}/{rem_output['id']}.png"
+        rem_output["screenshot"] = f"{rem_output['id']}.png"
     return rem_result

cs_storage/screenshot.py

@@ -15,8 +15,6 @@
     Template = None
     launch = None
 
-import cs_storage
-
 
 CURRENT_DIR = os.path.abspath(os.path.dirname(__file__))
 

cs_storage/tests/test_cs_storage.py

@@ -1,11 +1,8 @@
-import base64
 import io
 import json
-import uuid
-import zipfile
+import os
 
 import pytest
-import requests
 from marshmallow import exceptions
 
 import cs_storage
@@ -208,16 +205,24 @@ def test_cs_storage_serialization(exp_loc_res):
 def test_add_screenshot_links():
     rem_res = {"renderable": {"outputs": [{"id": "1234"}, {"id": "4567"}]}}
 
-    url = f"https://storage.googleapis.com/{cs_storage.BUCKET}/"
     assert cs_storage.add_screenshot_links(rem_res) == {
         "renderable": {
             "outputs": [
-                {"id": "1234", "screenshot": url + "1234.png"},
-                {"id": "4567", "screenshot": url + "4567.png"},
+                {"id": "1234", "screenshot": "1234.png"},
+                {"id": "4567", "screenshot": "4567.png"},
             ]
         }
     }
 
+    current_dir = os.path.abspath(os.path.dirname(__file__))
+    with open(f"{current_dir}/test-tc-remote.json") as f:
+        remote_outputs = json.loads(f.read())["outputs"]
+
+    cs_storage.add_screenshot_links(remote_outputs)
+    for output in remote_outputs["renderable"]["outputs"]:
+        link = output["screenshot"]
+        assert isinstance(cs_storage.read_screenshot(link), bytes)
+
 
 def test_errors():
     with pytest.raises(exceptions.ValidationError):

setup.py

@@ -17,7 +17,7 @@
     long_description_content_type="text/markdown",
     url="https://github.com/compute-tooling/compute-studio-storage",
     packages=setuptools.find_packages(),
-    install_requires=["marshmallow>=3.0.0", "gcsfs"],
+    install_requires=["marshmallow>=3.0.0", "fsspec", "gcsfs"],
     include_package_data=True,
     classifiers=[
         "Programming Language :: Python :: 3",