Merge pull request #41 from getmetal/task/sp/add-entity-metadata

Allow `metadata` for data entities

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "metal_sdk"
-version = "2.1.4"
+version = "2.2.0"
 authors = [
   { name="Metal Technologies Inc", email="[email protected]" },
 ]

src/metal_sdk/metal.py

@@ -308,20 +308,31 @@ def update_datasource(self, id: str, payload: dict):
         res = self.fetch("put", url, payload)
         return res
 
-    def __add_data_entity_resource(self, datasource, filename, file_size):
+    def __validate_metadata(self, metadata):
+        if metadata is not None:
+            if not isinstance(metadata, dict):
+                raise TypeError("Metadata must be a dictionary.")
+            for key, value in metadata.items():
+                if not isinstance(key, str) or not (isinstance(value, str) or isinstance(value, int)):
+                    raise TypeError("Metadata keys must be strings, and values must be strings or numbers.")
+
+    def __add_data_entity_resource(self, datasource, filename, file_size, metadata=None):
+        self.__validate_metadata(metadata)
         url = '/v1/data-entities'
         payload = {
             'datasource': datasource,
             'name': self.__sanitize_filename(filename),
             'sourceType': "file",
-            "status": "active"
         }
+
+        if metadata is not None:
+            payload['metadata'] = metadata
+
         headers = {'x-metal-file-size': str(file_size)}
 
         return self.fetch("post", url, payload, headers=headers)
 
-    def add_data_entity(self, datasource, file_path):
-
+    def add_data_entity(self, datasource, file_path, metadata=None):
         if datasource is None:
             raise ValueError("Payload must contain a 'datasource' id")
 
@@ -332,7 +343,7 @@ def add_data_entity(self, datasource, file_path):
         filename = os.path.basename(file_path)
         file_type, _ = mimetypes.guess_type(file_path)
 
-        resource = self.__add_data_entity_resource(datasource, filename, file_size)
+        resource = self.__add_data_entity_resource(datasource, filename, file_size, metadata=metadata)
         if not resource or 'data' not in resource:
             logger.error("Failed to create a data entity resource.")
             return None

src/metal_sdk/metal_async.py

@@ -309,19 +309,28 @@ async def update_datasource(self, id: str, payload: dict):
         res = await self.fetch("put", url, payload)
         return res
 
-    async def __add_data_entity_resource(self, datasource, filename, file_size):
+    def __validate_metadata(self, metadata):
+        if metadata is not None:
+            if not isinstance(metadata, dict):
+                raise TypeError("Metadata must be a dictionary.")
+            for key, value in metadata.items():
+                if not isinstance(key, str) or not (isinstance(value, str) or isinstance(value, int)):
+                    raise TypeError("Metadata keys must be strings, and values must be strings or numbers.")
+
+    async def __add_data_entity_resource(self, datasource, filename, file_size, metadata=None):
+        self.__validate_metadata(metadata)
         url = '/v1/data-entities'
         payload = {
             'datasource': datasource,
             'name': self.__sanitize_filename(filename),
             'sourceType': "file",
-            "status": "active"
+            "metadata": metadata
         }
         headers = {'x-metal-file-size': str(file_size)}
 
         return await self.fetch("post", url, payload, headers=headers)
 
-    async def add_data_entity(self, datasource, file_path):
+    async def add_data_entity(self, datasource, file_path, metadata=None):
         if datasource is None:
             raise ValueError("Payload must contain a 'datasource' id")
 
@@ -332,7 +341,7 @@ async def add_data_entity(self, datasource, file_path):
         filename = os.path.basename(file_path)
         file_type, _ = mimetypes.guess_type(file_path)
 
-        resource = await self.__add_data_entity_resource(datasource, filename, file_size)
+        resource = await self.__add_data_entity_resource(datasource, filename, file_size, metadata=metadata)
         if not resource or 'data' not in resource:
             logger.error("Failed to create a data entity resource.")
             return None

tests/test_metal.py

@@ -417,14 +417,17 @@ def test_add_data_entity(self):
             })
         metal._Metal__upload_file_to_url = mock.MagicMock()
 
-        res = metal.add_data_entity(datasource_id, mock_file_path)
+        res = metal.add_data_entity(datasource_id, mock_file_path, metadata={'vocalist': 'ozzy'})
 
         self.assertEqual(res['data']['url'], 'https://mockuploadurl.com')
 
         self.assertEqual(metal._Metal__add_data_entity_resource.call_count, 1)
         self.assertEqual(metal._Metal__upload_file_to_url.call_count, 1)
 
         create_args = metal._Metal__add_data_entity_resource.call_args[0]
+        create_kwargs = metal._Metal__add_data_entity_resource.call_args[1]
+        self.assertEqual(create_kwargs['metadata']['vocalist'], 'ozzy')
+
         self.assertEqual(create_args[0], datasource_id)
         self.assertEqual(create_args[1], "sample.csv")
         self.assertEqual(create_args[2], 47)