Add Asset.storage_type when retrieving assets from entitycore

src/entitysdk/models/asset.py

@@ -28,6 +28,13 @@ class Asset(Identifiable):
             description="The full s3 path of the asset.",
         ),
     ]
+    storage_type: Annotated[
+        StorageType,
+        Field(
+            examples=["aws_s3_open"],
+            description="Storage where the asset is located.",
+        ),
+    ]
     is_directory: Annotated[
         bool,
         Field(

tests/unit/downloaders/test_cell_morphology.py

@@ -19,6 +19,7 @@ def _mock_asset_response(asset_id):
         "status": "created",
         "meta": {},
         "sha256_digest": "sha256_digest",
+        "storage_type": "aws_s3_internal",
     }
 
 

tests/unit/downloaders/test_em_cell_mesh.py

@@ -17,6 +17,7 @@ def _mock_asset_response(asset_id):
         "meta": {},
         "sha256_digest": "sha256_digest",
         "label": "cell_surface_mesh",
+        "storage_type": "aws_s3_internal",
     }
 
 

tests/unit/downloaders/test_emodel.py

@@ -17,6 +17,7 @@ def _mock_asset_response(asset_id):
         "status": "created",
         "meta": {},
         "sha256_digest": "sha256_digest",
+        "storage_type": "aws_s3_internal",
     }
 
 

tests/unit/downloaders/test_ion_channel_model.py

@@ -17,6 +17,7 @@ def _mock_asset_response(asset_id):
         "status": "created",
         "meta": {},
         "sha256_digest": "sha256_digest",
+        "storage_type": "aws_s3_internal",
     }
 
 

tests/unit/downloaders/test_memodel.py

@@ -24,6 +24,7 @@ def _mock_morph_asset_response(asset_id):
         "status": "created",
         "meta": {},
         "sha256_digest": "sha256_digest",
+        "storage_type": "aws_s3_internal",
     }
 
 
@@ -40,6 +41,7 @@ def _mock_ic_asset_response(asset_id):
         "status": "created",
         "meta": {},
         "sha256_digest": "sha256_digest",
+        "storage_type": "aws_s3_internal",
     }
 
 
@@ -56,6 +58,7 @@ def _mock_emodel_asset_response(asset_id):
         "status": "created",
         "meta": {},
         "sha256_digest": "sha256_digest",
+        "storage_type": "aws_s3_internal",
     }
 
 

tests/unit/models/data/cell_morphology.json

@@ -45,6 +45,7 @@
       "meta": {},
       "label": "morphology",
       "id": "8c73ebe9-fe8a-4621-bc83-1f4dfc9bf5df",
+      "storage_type": "aws_s3_internal",
       "status": "created"
     },
     {
@@ -57,6 +58,7 @@
       "meta": {},
       "label": "morphology",
       "id": "4f63feb5-ee64-489b-9af3-59f3654c769c",
+      "storage_type": "aws_s3_internal",
       "status": "created"
     }
   ],

tests/unit/models/data/electrical_cell_recording.json

@@ -44,6 +44,7 @@
       "content_type": "application/nwb",
       "meta": {},
       "label": "nwb",
+      "storage_type": "aws_s3_internal",
       "id": "898ef0f4-d36f-4426-8af7-145cdb202004",
       "status": "created"
     }

tests/unit/models/data/ion_channel_model.json

@@ -1,19 +1,5 @@
 {
   "contributions": [],
-  "assets": [
-    {
-      "size": 1188,
-      "sha256_digest": "e80d50df8f069863e4dd1f361a8a049933402ed233342f4ddaafab87064943fb",
-      "path": "Ca_HVA2.mod",
-      "full_path": "public/a98b7abc-fc46-4700-9e3d-37137812c730/0dbced5f-cc3d-488a-8c7f-cfb8ea039dc6/assets/ion_channel_model/21ebb7ab-b41b-441d-b494-6665075d26b0/Ca_HVA2.mod",
-      "is_directory": false,
-      "content_type": "application/mod",
-      "meta": {},
-      "label": "neuron_mechanisms",
-      "id": "73e5c7ef-5cb6-484d-b534-2c1209f9f280",
-      "status": "created"
-    }
-  ],
   "license": null,
   "creation_date": "2022-10-20T13:05:35.588000Z",
   "update_date": "2025-04-15T12:41:30.454420Z",
@@ -99,6 +85,23 @@
         "main_ion": true
       }
     ],
-    "nonspecific": []
-  }
+    "nonspecific": [{"ihcn": "mA/cm2"}]
+  },
+  "assets": [
+    {
+      "id": "a90bab5c-8529-4829-99fc-e90a270e8b30",
+      "path": "Ca_HVA.mod",
+      "full_path": "public/a98b7abc-fc46-4700-9e3d-37137812c730/0dbced5f-cc3d-488a-8c7f-cfb8ea039dc6/assets/ion_channel_model/9720aa4a-0cd1-4f10-b6ad-04deb8f997b6/Ca_HVA.mod",
+      "is_directory": false,
+      "content_type": "application/mod",
+      "label": "neuron_mechanisms",
+      "storage_type": "aws_s3_internal",
+      "size": 561221,
+      "sha256_digest": "8b850eface3e7b881db62c361dc5da62be008ca81501c406a59646bf7075c95b",
+      "status": "created",
+      "creation_date": null,
+      "update_date": null,
+      "meta": {}
+    }
+  ]
 }

tests/unit/models/data/ion_channel_recording.json

@@ -45,6 +45,7 @@
       "meta": {},
       "label": "nwb",
       "id": "898ef0f4-d36f-4426-8af7-145cdb202004",
+      "storage_type": "aws_s3_internal",
       "status": "created"
     }
   ],

tests/unit/models/test_asset.py

@@ -9,6 +9,7 @@ def test_asset():
         id=MOCK_UUID,
         path="path/to/asset",
         full_path="full/path/to/asset",
+        storage_type=StorageType.aws_s3_internal,
         label=AssetLabel.sonata_circuit,
         is_directory=False,
         content_type=ContentType.text_plain,
@@ -28,6 +29,7 @@ def test_asset():
         "sha256_digest": None,
         "meta": {},
         "label": AssetLabel.sonata_circuit,
+        "storage_type": StorageType.aws_s3_internal,
     }
 
 

tests/unit/staging/conftest.py

@@ -13,6 +13,7 @@
     Species,
     Subject,
 )
+from entitysdk.types import StorageType
 
 DATA_DIR = Path(__file__).parent / "data"
 
@@ -65,6 +66,7 @@ def circuit(subject, brain_region):
                 path="circuit",
                 full_path="/circuit",
                 is_directory=True,
+                storage_type=StorageType.aws_s3_internal,
             )
         ],
     )
@@ -143,6 +145,7 @@ def simulation(circuit):
                 full_path="/foo.json",
                 size=0,
                 is_directory=False,
+                storage_type=StorageType.aws_s3_internal,
             ),
             Asset(
                 id=uuid.uuid4(),
@@ -152,6 +155,7 @@ def simulation(circuit):
                 full_path="/bar.json",
                 size=0,
                 is_directory=False,
+                storage_type=StorageType.aws_s3_internal,
             ),
             Asset(
                 id=uuid.uuid4(),
@@ -161,6 +165,7 @@ def simulation(circuit):
                 full_path="/PoissonInputStimulus_spikes_1.h5",
                 size=0,
                 is_directory=False,
+                storage_type=StorageType.aws_s3_internal,
             ),
             Asset(
                 id=uuid.uuid4(),
@@ -170,6 +175,7 @@ def simulation(circuit):
                 full_path="/PoissonInputStimulus_spikes_2.h5",
                 size=0,
                 is_directory=False,
+                storage_type=StorageType.aws_s3_internal,
             ),
         ],
     )
@@ -232,6 +238,7 @@ def simulation_result(simulation):
                 full_path="/soma_voltage1.h5",
                 size=0,
                 is_directory=False,
+                storage_type=StorageType.aws_s3_internal,
             ),
             Asset(
                 id=uuid.uuid4(),
@@ -241,6 +248,7 @@ def simulation_result(simulation):
                 full_path="/soma_voltage2.h5",
                 size=0,
                 is_directory=False,
+                storage_type=StorageType.aws_s3_internal,
             ),
             Asset(
                 id=uuid.uuid4(),
@@ -250,6 +258,7 @@ def simulation_result(simulation):
                 full_path="/out.h5",
                 size=0,
                 is_directory=False,
+                storage_type=StorageType.aws_s3_internal,
             ),
         ],
     )

tests/unit/staging/data/memodel.json

@@ -121,6 +121,7 @@
         "sha256_digest": "f6f5c9cbf6472411bd34f24fcf960ddcebc61aa311f27e2472f75544b99cbe06",
         "status": "created",
         "meta": {},
+        "storage_type": "aws_s3_internal",
         "label": "morphology"
       },
       {
@@ -133,6 +134,7 @@
         "sha256_digest": "f42d6a64f9a554d5fd7e52c34643073ef4d28dfebb5a3acd7f470c6f1b014b9d",
         "status": "created",
         "meta": {},
+        "storage_type": "aws_s3_internal",
         "label": "morphology"
       },
       {
@@ -145,6 +147,7 @@
         "sha256_digest": "44a476456bd6db40eea5ef9a6202a708a37c9d6aeba5f3edb9f8ba6e3eec10b6",
         "status": "created",
         "meta": {},
+        "storage_type": "aws_s3_internal",
         "label": "morphology"
       }
     ],
@@ -340,6 +343,7 @@
         "sha256_digest": "69b206cfdfacdba2b406b6fc8a7e305577d0bc407e6edcd128e703f7081c2eed",
         "status": "created",
         "meta": {},
+        "storage_type": "aws_s3_internal",
         "label": "neuron_hoc"
       },
       {
@@ -352,6 +356,7 @@
         "sha256_digest": "5d9b58544d7b82b82d07e97e1a5698a488a4b1bf5f49d1984ba72b1ce280abe1",
         "status": "created",
         "meta": {},
+        "storage_type": "aws_s3_internal",
         "label": "emodel_optimization_output"
       }
     ],

tests/unit/test_client.py

@@ -148,6 +148,7 @@ def _mock_asset_response(
         "meta": {},
         "sha256_digest": "sha256_digest",
         "label": label,
+        "storage_type": "aws_s3_internal",
     }
 
 
@@ -548,6 +549,7 @@ def _mock_asset_delete_response(asset_id):
         "meta": {},
         "id": str(asset_id),
         "status": "deleted",
+        "storage_type": "aws_s3_internal",
     }
 
 
@@ -761,6 +763,7 @@ def test_client_download_assets__entity(
                 id=asset_id,
                 path="foo.json",
                 full_path="/foo/asset1",
+                storage_type=StorageType.aws_s3_internal,
                 is_directory=False,
                 content_type="application/json",
                 label="cell_composition_summary",
@@ -832,6 +835,7 @@ def test_upload_directory_by_paths(
         "sha256_digest": None,
         "size": -1,
         "status": "created",
+        "storage_type": "aws_s3_internal",
     }
     httpx_mock.add_response(
         method="POST",
@@ -1096,6 +1100,7 @@ def test_client_download_directory__asset(
         id=asset_id,
         path="path_to_asset",
         full_path="/circuit",
+        storage_type=StorageType.aws_s3_internal,
         is_directory=True,
         size=0,
         content_type="application/vnd.directory",

tests/unit/utils/test_asset.py

@@ -4,7 +4,7 @@
 
 from entitysdk.exception import EntitySDKError
 from entitysdk.models import Asset
-from entitysdk.types import AssetLabel, ContentType
+from entitysdk.types import AssetLabel, ContentType, StorageType
 from entitysdk.utils import asset as test_module
 
 
@@ -15,6 +15,7 @@ def assets():
             id=uuid.uuid4(),
             path="foo/asset1",
             full_path="/foo/asset1",
+            storage_type=StorageType.aws_s3_internal,
             is_directory=False,
             content_type=ContentType.application_json,
             size=1,
@@ -24,6 +25,7 @@ def assets():
             id=uuid.uuid4(),
             path="foo/asset2",
             full_path="/foo/asset2",
+            storage_type=StorageType.aws_s3_internal,
             is_directory=False,
             content_type=ContentType.text_plain,
             size=1,
@@ -33,6 +35,7 @@ def assets():
             id=uuid.uuid4(),
             path="foo/asset3",
             full_path="/foo/asset3",
+            storage_type=StorageType.aws_s3_internal,
             is_directory=False,
             content_type=ContentType.text_plain,
             size=1,