added support for fetching cas values from legacy storage

Author: pgm

breadbox/breadbox/api/dependencies.py

@@ -18,6 +18,11 @@ def get_db_with_user(request: Request):
         db.close()
 
 
+def get_legacy_cas_bucket():
+    settings = get_settings()
+    return settings.LEGACY_CAS_BUCKET
+
+
 def get_cas_db_path():
     settings = get_settings()
     return os.path.join(settings.filestore_location, "cas.sqlite3")

breadbox/breadbox/api/temp/cas.py

@@ -1,17 +1,28 @@
 from .router import router
 from breadbox.schemas.cas import CASKey, CASValue
 from breadbox.io import cas
-import os.path
-from typing import Annotated
+from typing import Annotated, Optional
 from fastapi import APIRouter, Body, Depends, HTTPException
-from breadbox.api.dependencies import get_db_with_user, get_cas_db_path
+from breadbox.api.dependencies import (
+    get_db_with_user,
+    get_cas_db_path,
+    get_legacy_cas_bucket,
+)
 
 # Methods for getting/setting values in Content-addressable-storage (CAS)
 @router.get(
     "/cas/{key}", operation_id="get_cas_value", response_model=CASValue,
 )
-def get_cas_value(key: str, cas_db_path: Annotated[str, Depends(get_cas_db_path)]):
+def get_cas_value(
+    key: str,
+    cas_db_path: Annotated[str, Depends(get_cas_db_path)],
+    legacy_cas_bucket: Annotated[Optional[str], Depends(get_legacy_cas_bucket)],
+):
     value = cas.get_value(cas_db_path, key)
+
+    if legacy_cas_bucket is not None and value is None:
+        value = cas.legacy_get_value(legacy_cas_bucket, key)
+
     if value is None:
         raise HTTPException(status_code=404)
     return CASValue(value=value)
@@ -21,7 +32,14 @@ def get_cas_value(key: str, cas_db_path: Annotated[str, Depends(get_cas_db_path)
     "/cas", operation_id="set_cas_value", response_model=CASKey,
 )
 def set_cas_value(
-    value: CASValue, cas_db_path: Annotated[str, Depends(get_cas_db_path)]
+    value: CASValue,
+    cas_db_path: Annotated[str, Depends(get_cas_db_path)],
+    legacy_cas_bucket: Annotated[Optional[str], Depends(get_legacy_cas_bucket)],
 ):
     key = cas.set_value(cas_db_path, value.value)
+
+    if legacy_cas_bucket is not None:
+        legacy_key = cas.legacy_set_value(legacy_cas_bucket, value.value)
+        assert legacy_key == key
+
     return CASKey(key=key)

breadbox/breadbox/config.py

@@ -29,6 +29,8 @@ class Settings(BaseSettings):
         "CELERY_RESULT_BACKEND", "redis://127.0.0.1:6379/0"
     )
 
+    LEGACY_CAS_BUCKET: Optional[str] = os.environ.get("LEGACY_CAS_BUCKET")
+
     model_config = SettingsConfigDict(env_file=".env",)
 
     @field_validator("host_scheme_override")

breadbox/breadbox/io/cas.py

@@ -3,6 +3,7 @@
 from base64 import urlsafe_b64encode
 from sqlitedict import SqliteDict
 from typing import Optional
+from google.cloud import storage
 
 
 def get_value(db_path: str, key: str) -> Optional[str]:
@@ -20,3 +21,35 @@ def set_value(db_path: str, value: str) -> str:
         db[key] = value_bytes
         db.commit()
     return key
+
+
+##########
+# Legacy methods for getting/setting values in google bucket
+
+
+def _get_private_datasets_bucket(bucket_name: str) -> storage.Bucket:
+    client = storage.Client()
+
+    bucket = client.bucket(bucket_name)
+    return bucket
+
+
+def legacy_get_value(bucket_name: str, key: str) -> Optional[str]:
+    bucket = _get_private_datasets_bucket(bucket_name)
+
+    blob = bucket.blob(key)
+    if not blob.exists():
+        return None
+
+    return blob.download_as_string().decode("utf8")
+
+
+def legacy_set_value(bucket_name: str, value: str) -> Optional[str]:
+    bucket = _get_private_datasets_bucket(bucket_name)
+
+    value_bytes = value.encode("utf8")
+    key = urlsafe_b64encode(hashlib.sha256(value_bytes).digest()).decode("utf8")
+    blob = bucket.blob(key)
+    if not blob.exists():
+        blob.upload_from_string(value_bytes)
+    return key

breadbox/poetry.lock

@@ -33,6 +33,7 @@ depmap-compute = {version = "0.1.13", source = "public-python"}
 # depmap-compute = {path = "../depmap-compute", develop = true}
 fastparquet = "^2024.5.0"
 sqlitedict="^2.1.0"
+google-cloud-storage = "^3.1.0"
 
 [tool.poetry.group.dev.dependencies]
 pyright-ratchet = {git = "https://github.com/pgm/pyright-ratchet.git", rev = "v0.3.1"}