Taiga client v3 updates

taigapy/client_v3.py

@@ -1,6 +1,8 @@
 import pandas as pd
 from typing import Optional, Dict, List, Any
 import os
+
+from taigapy.custom_exceptions import Taiga404Exception, TaigaHttpException
 from .taiga_api import TaigaApi
 import re
 from enum import Enum, auto
@@ -10,7 +12,7 @@
 import colorful as cf
 import uuid
 from typing import Callable, Tuple
-from .types import DatasetVersionMetadataDict
+from .types import DatasetMetadataDict, DatasetVersionMetadataDict
 from .simple_cache import Cache
 from .types import DataFileUploadFormat
 from .format_utils import read_hdf5, read_parquet, convert_csv_to_parquet
@@ -28,6 +30,7 @@ class LocalFormat(Enum):
     PARQUET_TABLE = "parquet_table"
     CSV_TABLE = "csv_table"
     CSV_MATRIX = "csv_matrix"
+    RAW = "raw"
 
 
 # the different formats files might be stored in on Taiga.
@@ -311,6 +314,31 @@ def _get(self, datafile_id: str) -> pd.DataFrame:
             )
 
         return result
+
+    def upload_to_gcs(self, data_file_taiga_id: str, dest_gcs_path: str) -> bool:
+        """Upload a Taiga datafile to a specified location in Google Cloud Storage.
+
+        The service account [email protected] must have
+        storage.buckets.create access for this request.
+
+        Arguments:
+            `data_file_taiga_id` -- Taiga ID in the form dataset_permaname.dataset_version/datafile_name
+            `dest_gcs_path` -- Google Storage path to upload to, in the form bucket:path
+
+        Returns:
+            bool -- Whether the file was successfully uploaded
+        """
+        full_taiga_id = self.get_canonical_id(data_file_taiga_id)
+        if full_taiga_id is None:
+            return False
+
+        try:
+            self.api.upload_to_gcs(full_taiga_id, dest_gcs_path)
+            return True
+        except (ValueError, TaigaHttpException) as e:
+            print(cf.red(str(e)))
+            return False
+
 
     def download_to_cache(self, datafile_id: str, requested_format: Union[LocalFormat, str]) -> str:
         """
@@ -372,6 +400,13 @@ def download_to_cache(self, datafile_id: str, requested_format: Union[LocalForma
                 raise Exception(
                     f"Requested {requested_format} but taiga_format={taiga_format}"
                 )
+        elif requested_format == LocalFormat.RAW:
+            if taiga_format == TaigaStorageFormat.RAW_BYTES:
+                local_path = self._download_to_cache(canonical_id)
+            else:
+                raise Exception(
+                    f"Requested {requested_format} but taiga_format={taiga_format}"
+                )
         else:
             raise Exception(
                 f"Requested {requested_format} but taiga_format={taiga_format}"
@@ -440,9 +475,10 @@ def _upload_taiga_reference(file: TaigaReference):
             self.api.upload_file_to_taiga(
                 upload_session_id,
                 {
-                    "filename": upload.name,
+                    "filename": file.name,
                     "filetype": "virtual",
-                    "existingTaigaId": upload.taiga_id,
+                    "existingTaigaId": file.taiga_id,
+                    "custom_metadata": file.custom_metadata,
                 },
             )
 
@@ -498,6 +534,30 @@ def _dataset_version_summary(self, dataset_version_id):
                 for f in version_metadata["datasetVersion"]["datafiles"]
             ],
         )
+
+    def get_dataset_metadata(
+            self,
+            permaname: str,
+            version: Optional[str] = None
+        ) -> Optional[Union[DatasetMetadataDict, DatasetVersionMetadataDict]]:
+        """Get metadata about a dataset
+
+        Keyword Arguments:
+            - `permaname` -- Datafile ID of the datafile to get, in the form dataset_permaname
+            - `dataset_version`: Either the numerical version (if `dataset_permaname` is provided) 
+            or the unique Taiga dataset version id
+
+        Returns:
+            Union[DatasetMetadataDict, DatasetVersionMetadataDict]
+            `DatasetMetadataDict` if only permaname provided.  `DatasetVersionMetadataDict` if both permaname and version provided.
+        """
+        try:
+            return self.api.get_dataset_version_metadata(
+                dataset_permaname=permaname, dataset_version=version
+            )
+        except (ValueError, Taiga404Exception) as e:
+            print(cf.red(str(e)))
+            return None
 
     def create_dataset(
         self,
@@ -628,10 +688,7 @@ def _get_taiga_storage_format(self, datafile_id: str) -> TaigaStorageFormat:
         if metadata.type == "Columnar":
             return TaigaStorageFormat.CSV_TABLE
         if metadata.type == "Raw":
-            value = metadata.custom_metadata.get(
-                "client_storage_format", TaigaStorageFormat.RAW_BYTES
-            )
-            return TaigaStorageFormat(value)
+            return TaigaStorageFormat.RAW_BYTES
         else:
             raise Exception(f"unknown type: {metadata.type}")
 

taigapy/taiga_api.py

@@ -305,6 +305,23 @@ def get_datafile_metadata(
     def get_dataset_version_metadata(
         self, dataset_permaname: str, dataset_version: Optional[str]
     ) -> Union[DatasetMetadataDict, DatasetVersionMetadataDict]:
+        """
+        Get metadata about a dataset
+
+        Keyword Arguments:
+        - `dataset_permaname`: The Taiga dataset permaname
+        - `dataset_version`: Either the numerical version (if `dataset_permaname` is provided) or the unique Taiga dataset version id (if `dataset_permaname` is not provided)
+
+        Returns:
+            Union[DatasetMetadataDict, DatasetVersionMetadataDict] -- See docs at https://github.com/broadinstitute/taigapy for more details
+            `DatasetMetadataDict` if only `dataset_permaname` param is provided. Returns `DatasetVersionMetadataDict` if `dataset_version` is provided.
+
+        A Taiga dataset id consists of:
+            - `permaname`: Unique identifier for dataset group
+            - `version`: Numerical version of the Taiga dataset upload
+            - `name`: Name of the file in the Taiga dataset group. If provided the dataset id is for that specific file rather than dataset group.
+
+        """
         api_endpoint = "/api/dataset/{}".format(dataset_permaname)
         if dataset_version is not None:
             api_endpoint = "{}/{}".format(api_endpoint, dataset_version)