qwak-core 0.1.37__py3-none-any.whl → 0.5.19__py3-none-any.whl

qwak/clients/feature_store/management_client.py

@@ -16,28 +16,37 @@ from _qwak_proto.qwak.feature_store.entities.entity_service_pb2_grpc import (
 )
 from _qwak_proto.qwak.feature_store.features.feature_set_pb2 import FeatureSetSpec
 from _qwak_proto.qwak.feature_store.features.feature_set_service_pb2 import (
-    CreateUploadURLRequest,
     CreateUploadURLRequestV1,
     DeleteFeatureSetRequest,
     DeleteFeatureSetResponse,
+    DeleteFeaturesetVersionRequest,
+    GetEnvToFeatureSetsMappingRequest,
+    GetEnvToFeatureSetsMappingResponse,
     GetFeatureSetByNameRequest,
     GetFeatureSetByNameResponse,
     GetFeaturesetSchedulingStateRequest,
     GetFeaturesetSchedulingStateResponse,
     ListFeatureSetsRequest,
     ListFeatureSetsResponse,
+    ListFeaturesetVersionsByNameRequest,
+    ListFeaturesetVersionsByNameResponse,
     PauseFeaturesetRequest,
     RegisterFeatureSetRequest,
     RegisterFeatureSetResponse,
     ResumeFeaturesetRequest,
     RunBatchFeatureSetRequest,
+    RunBatchFeatureSetResponse,
     RunParquetBatchFeatureSetRequest,
+    SetActiveFeaturesetVersionRequest,
     UpdateFeatureSetRequest,
     UpdateFeatureSetResponse,
 )
 from _qwak_proto.qwak.feature_store.features.feature_set_service_pb2_grpc import (
     FeatureSetServiceStub,
 )
+from _qwak_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    Transformation as ProtoTransformation,
+)
 from _qwak_proto.qwak.feature_store.sources.data_source_service_pb2 import (
     CreateDataSourceRequest,
     CreateDataSourceResponse,
@@ -54,10 +63,18 @@ from _qwak_proto.qwak.feature_store.sources.data_source_service_pb2 import (
 from _qwak_proto.qwak.feature_store.sources.data_source_service_pb2_grpc import (
     DataSourceServiceStub,
 )
+from _qwak_proto.qwak.feature_store.v1.internal.data_source.data_source_service_pb2 import (
+    GetDataSourceSourceCodeUploadRequest,
+    GetDataSourceSourceCodeUploadResponse,
+)
+from _qwak_proto.qwak.feature_store.v1.internal.data_source.data_source_service_pb2_grpc import (
+    DataSourceServiceStub as InternalDataSourceServiceStub,
+)
 from dependency_injector.wiring import Provide
 from grpc import RpcError, StatusCode
 from qwak.exceptions import QwakException
 from qwak.inner.di_configuration import QwakContainer
+from qwak.inner.tool.grpc.grpc_try_wrapping import grpc_try_catch_wrapper
 
 
 class FeatureRegistryClient:
@@ -69,16 +86,31 @@ class FeatureRegistryClient:
         self._features_service = FeatureSetServiceStub(grpc_channel)
         self._entity_service = EntityServiceStub(grpc_channel)
         self._sources_service = DataSourceServiceStub(grpc_channel)
+        self._internal_data_sources_service = InternalDataSourceServiceStub(
+            grpc_channel
+        )
+
+    @grpc_try_catch_wrapper(
+        "Failed to get a presigned url", reraise_non_grpc_error_original_exception=True
+    )
+    def get_datasource_source_code_presign_url(
+        self, ds_name: str, repository_name: str
+    ) -> GetDataSourceSourceCodeUploadResponse:
+        return self._internal_data_sources_service.GetDataSourceSourceCodeUploadURL(
+            GetDataSourceSourceCodeUploadRequest(
+                data_source_name=ds_name, repository_name=repository_name
+            )
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to list features", reraise_non_grpc_error_original_exception=True
+    )
     def list_feature_sets(self) -> ListFeatureSetsResponse:
         """
         Returns:
             ListFeatureSetsResponse
         """
-        try:
-            return self._features_service.ListFeatureSets(ListFeatureSetsRequest())
-        except RpcError as e:
-            raise QwakException(f"Failed to list features, error is {e.details()}")
+        return self._features_service.ListFeatureSets(ListFeatureSetsRequest())
 
     def get_entity_by_name(self, entity_name: str) -> GetEntityByNameResponse:
         """
@@ -95,6 +127,9 @@ class FeatureRegistryClient:
             "entity",
         )
 
+    @grpc_try_catch_wrapper(
+        "Failed to delete feature", reraise_non_grpc_error_original_exception=True
+    )
     def delete_feature_set(self, feature_set_id: str) -> DeleteFeatureSetResponse:
         """
         Args:
@@ -102,13 +137,13 @@ class FeatureRegistryClient:
         Returns:
             DeleteFeatureSetResponse
         """
-        try:
-            return self._features_service.DeleteFeatureSet(
-                DeleteFeatureSetRequest(feature_set_id=feature_set_id)
-            )
-        except RpcError as e:
-            raise QwakException(f"Failed to delete feature, error is {e.details()}")
+        return self._features_service.DeleteFeatureSet(
+            DeleteFeatureSetRequest(feature_set_id=feature_set_id)
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to delete entity", reraise_non_grpc_error_original_exception=True
+    )
     def delete_entity(self, entity_id: str) -> DeleteEntityResponse:
         """
         Args:
@@ -116,13 +151,13 @@ class FeatureRegistryClient:
         Returns:
             DeleteEntityResponse
         """
-        try:
-            return self._entity_service.DeleteEntity(
-                DeleteEntityRequest(entity_id=entity_id)
-            )
-        except RpcError as e:
-            raise QwakException(f"Failed to delete entity, error is {e.details()}")
+        return self._entity_service.DeleteEntity(
+            DeleteEntityRequest(entity_id=entity_id)
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to delete data source", reraise_non_grpc_error_original_exception=True
+    )
     def delete_data_source(self, data_source_id: str) -> DeleteDataSourceResponse:
         """
         Args:
@@ -130,28 +165,28 @@ class FeatureRegistryClient:
         Returns:
             DeleteDataSourceResponse
         """
-        try:
-            return self._sources_service.DeleteDataSource(
-                DeleteDataSourceRequest(data_source_id=data_source_id)
-            )
-        except RpcError as e:
-            raise QwakException(f"Failed to delete data source, error is {e.details()}")
+        return self._sources_service.DeleteDataSource(
+            DeleteDataSourceRequest(data_source_id=data_source_id)
+        )
 
-    def run_feature_set(self, feature_set_name: str):
+    @grpc_try_catch_wrapper(
+        "Failed to run feature", reraise_non_grpc_error_original_exception=True
+    )
+    def run_feature_set(self, feature_set_name: str) -> RunBatchFeatureSetResponse:
         """
         Trigger a single ETL run for that feature set name
         Args:
             feature_set_name: feature set name to run
         Returns:
-            None
+            :return RunBatchFeatureSetResponse with the execution id
         """
-        try:
-            self._features_service.RunBatchFeatureSet(
-                RunBatchFeatureSetRequest(feature_set_name=feature_set_name)
-            )
-        except RpcError as e:
-            raise QwakException(f"Failed to run feature, error is {e.details()}")
+        return self._features_service.RunBatchFeatureSet(
+            RunBatchFeatureSetRequest(feature_set_name=feature_set_name)
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to run feature", reraise_non_grpc_error_original_exception=True
+    )
     def run_parquet_batch_feature_set(
         self, feature_set_name: str, parquet_location: str, timestamp_column: str
     ):
@@ -164,17 +199,17 @@ class FeatureRegistryClient:
         Returns:
             None
         """
-        try:
-            self._features_service.RunParquetBatchFeatureSet(
-                RunParquetBatchFeatureSetRequest(
-                    feature_set_name=feature_set_name,
-                    parquet_location=parquet_location,
-                    timestamp_column=timestamp_column,
-                )
+        self._features_service.RunParquetBatchFeatureSet(
+            RunParquetBatchFeatureSetRequest(
+                feature_set_name=feature_set_name,
+                parquet_location=parquet_location,
+                timestamp_column=timestamp_column,
             )
-        except RpcError as e:
-            raise QwakException(f"Failed to run feature, error is {e.details()}")
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to update entity", reraise_non_grpc_error_original_exception=True
+    )
     def update_entity(
         self, entity_id: str, proto_entity_spec: EntitySpec
     ) -> UpdateEntityResponse:
@@ -185,13 +220,13 @@ class FeatureRegistryClient:
         Returns:
             UpdateEntityResponse
         """
-        try:
-            return self._entity_service.UpdateEntity(
-                UpdateEntityRequest(entity_id=entity_id, entity_spec=proto_entity_spec)
-            )
-        except RpcError as e:
-            raise QwakException(f"Failed to update entity, error is {e.details()}")
+        return self._entity_service.UpdateEntity(
+            UpdateEntityRequest(entity_id=entity_id, entity_spec=proto_entity_spec)
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to create entity", reraise_non_grpc_error_original_exception=True
+    )
     def create_entity(self, proto_entity_spec: EntitySpec) -> CreateEntityResponse:
         """
         Args:
@@ -199,32 +234,29 @@ class FeatureRegistryClient:
         Returns:
             CreateEntityResponse
         """
-        try:
-            return self._entity_service.CreateEntity(
-                CreateEntityRequest(entity_spec=proto_entity_spec)
-            )
-        except RpcError as e:
-            raise QwakException(f"Failed to create entity, error is {e.details()}")
+        return self._entity_service.CreateEntity(
+            CreateEntityRequest(entity_spec=proto_entity_spec)
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to list entities", reraise_non_grpc_error_original_exception=True
+    )
     def list_entities(self) -> ListEntitiesResponse:
         """
         Returns:
              ListEntitiesResponse
         """
-        try:
-            return self._entity_service.ListEntities(ListEntitiesRequest())
-        except RpcError as e:
-            raise QwakException(f"Failed to list entities, error is {e.details()}")
+        return self._entity_service.ListEntities(ListEntitiesRequest())
 
+    @grpc_try_catch_wrapper(
+        "Failed to list data sources", reraise_non_grpc_error_original_exception=True
+    )
     def list_data_sources(self) -> ListDataSourcesResponse:
         """
         Returns:
             ListDataSourcesResponse
         """
-        try:
-            return self._sources_service.ListDataSources(ListDataSourcesRequest())
-        except RpcError as e:
-            raise QwakException(f"Failed to list data sources, error is {e.details()}")
+        return self._sources_service.ListDataSources(ListDataSourcesRequest())
 
     def get_data_source_by_name(
         self, data_source_name: str
@@ -243,6 +275,9 @@ class FeatureRegistryClient:
             "data source",
         )
 
+    @grpc_try_catch_wrapper(
+        "Failed to update data source", reraise_non_grpc_error_original_exception=True
+    )
     def update_data_source(
         self,
         data_source_id,
@@ -255,16 +290,16 @@ class FeatureRegistryClient:
         Returns:
              UpdateDataSourceResponse
         """
-        try:
-            return self._sources_service.UpdateDataSource(
-                UpdateDataSourceRequest(
-                    data_source_id=data_source_id,
-                    data_source_spec=proto_data_source,
-                )
+        return self._sources_service.UpdateDataSource(
+            UpdateDataSourceRequest(
+                data_source_id=data_source_id,
+                data_source_spec=proto_data_source,
             )
-        except RpcError as e:
-            raise QwakException(f"Failed to update data source, error is {e.details()}")
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to create data source", reraise_non_grpc_error_original_exception=True
+    )
     def create_data_source(
         self, proto_data_source: FeatureSetSpec
     ) -> CreateDataSourceResponse:
@@ -274,12 +309,9 @@ class FeatureRegistryClient:
         Returns:
             CreateDataSourceResponse
         """
-        try:
-            return self._sources_service.CreateDataSource(
-                CreateDataSourceRequest(data_source_spec=proto_data_source)
-            )
-        except RpcError as e:
-            raise QwakException(f"Failed to create data source, error is {e.details()}")
+        return self._sources_service.CreateDataSource(
+            CreateDataSourceRequest(data_source_spec=proto_data_source)
+        )
 
     def get_feature_set_by_name(self, feature_set_name) -> GetFeatureSetByNameResponse:
         """
@@ -296,6 +328,9 @@ class FeatureRegistryClient:
             "feature set",
         )
 
+    @grpc_try_catch_wrapper(
+        "Failed to update feature set", reraise_non_grpc_error_original_exception=True
+    )
     def update_feature_set(
         self,
         feature_set_id: str,
@@ -308,16 +343,15 @@ class FeatureRegistryClient:
         Returns:
             UpdateFeatureSetResponse
         """
-
-        try:
-            return self._features_service.UpdateFeatureSet(
-                UpdateFeatureSetRequest(
-                    feature_set_id=feature_set_id, feature_set_spec=proto_feature_set
-                )
+        return self._features_service.UpdateFeatureSet(
+            UpdateFeatureSetRequest(
+                feature_set_id=feature_set_id, feature_set_spec=proto_feature_set
             )
-        except RpcError as e:
-            raise QwakException(f"Failed to update feature set, error is {e.details()}")
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to create feature set", reraise_non_grpc_error_original_exception=True
+    )
     def create_feature_set(
         self, proto_feature_set: FeatureSetSpec
     ) -> RegisterFeatureSetResponse:
@@ -327,13 +361,9 @@ class FeatureRegistryClient:
         Returns:
             RegisterFeatureSetResponse
         """
-
-        try:
-            return self._features_service.RegisterFeatureSet(
-                RegisterFeatureSetRequest(feature_set_spec=proto_feature_set)
-            )
-        except RpcError as e:
-            raise QwakException(f"Failed to create feature set, error is {e.details()}")
+        return self._features_service.RegisterFeatureSet(
+            RegisterFeatureSetRequest(feature_set_spec=proto_feature_set)
+        )
 
     @staticmethod
     def _fetch_by_name(extract_function, name, fs_object_type):
@@ -355,30 +385,15 @@ class FeatureRegistryClient:
                     f"Failed to fetch {fs_object_type} by name [{name}], error code: [{e.details()}]"
                 )
 
-    def get_presigned_url(self, feature_set_name: str):
-        """
-        get presigned url according feature set name
-        Args:
-            feature_set_name: name of the feature set
-        Returns:
-            presigned url
-        """
-        try:
-            return self._features_service.CreateUploadURL(
-                CreateUploadURLRequest(feature_set_name=feature_set_name)
-            ).upload_url
-        except Exception as e:
-            raise QwakException(
-                f"Failed to get presigned url for feature name: {feature_set_name} , error code: [{e.details()}]"
-            )
-
-    def get_presigned_url_v1(self, feature_set_name: str, transformation) -> str:
+    def get_presigned_url_v1(
+        self, feature_set_name: str, transformation: ProtoTransformation
+    ) -> str:
         """
         Get presigned url according feature set name and transformation type
         @param feature_set_name: featureset name for which we're getting the presigned URL
         @type feature_set_name: str
-        @param transformation: The transformation for which we're getting the presigned URL
-        @type transformation: BatchTransformation
+        @param transformation: The proto object of the transformation for which we're getting the presigned URL
+        @type transformation: ProtoTransformation
         @return:
         @rtype:
         """
@@ -386,65 +401,107 @@ class FeatureRegistryClient:
             return self._features_service.CreateUploadURLV1(
                 CreateUploadURLRequestV1(
                     featureset_name=feature_set_name,
-                    transformation=transformation._to_proto(),
+                    transformation=transformation,
                 )
             ).upload_url
-        except Exception as e:
+        except RpcError as e:
             raise QwakException(
                 f"Failed to get presigned url for featureset name {feature_set_name} and transformation type "
                 f"{type(transformation)}, error code: [{e.details()}]"
             )
 
-    def get_data_source_presigned_url(
+    @grpc_try_catch_wrapper(
+        "Failed to get presigned url for data source: {data_source_name}, object name: {object_name}",
+        reraise_non_grpc_error_original_exception=True,
+    )
+    def get_data_source_presigned_url_v1(
         self, data_source_name: str, object_name: str
     ) -> str:
         """
-        create a pre-sign url for data source according to qwak path convention
+        create a pre-sign url for data source according to qwak path convention that includes a timestamp millis
+        to avoid using the same path. This method is meant to be used by the new sdk, which passes the path along
+        in the proto message.
         Args:
             data_source_name: str:  data source name
             object_name: str object name
         Returns:
             presigned url: str
         """
-        try:
-            return self._sources_service.CreateDataSourceUploadURL(
-                CreateDataSourceUploadURLRequest(
-                    data_source_name=data_source_name, object_name=object_name
-                )
-            ).upload_url
-        except Exception as e:
-            raise QwakException(
-                f"Failed to get presigned url for data source: {data_source_name}, object name: {object_name} , error code: [{e.details()}]"
-            ) from e
+        return self._sources_service.CreateDataSourceUploadURLV1(
+            CreateDataSourceUploadURLRequest(
+                data_source_name=data_source_name, object_name=object_name
+            )
+        ).upload_url
 
+    @grpc_try_catch_wrapper(
+        "Failed to query scheduling state for Featureset {feature_set_name}",
+        reraise_non_grpc_error_original_exception=True,
+    )
     def get_feature_set_scheduling_state(
         self, feature_set_name: str
     ) -> GetFeaturesetSchedulingStateResponse:
-        try:
-            return self._features_service.GetFeaturesetSchedulingState(
-                GetFeaturesetSchedulingStateRequest(featureset_name=feature_set_name)
-            )
-        except Exception as e:
-            raise QwakException(
-                f"Failed to query scheduling state for Featureset {feature_set_name}, nested exception: [{e.details()}]"
-            ) from e
+        return self._features_service.GetFeaturesetSchedulingState(
+            GetFeaturesetSchedulingStateRequest(featureset_name=feature_set_name)
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to pause Featureset {feature_set_name}",
+        reraise_non_grpc_error_original_exception=True,
+    )
     def pause_feature_set(self, feature_set_name: str):
-        try:
-            self._features_service.PauseFeatureset(
-                PauseFeaturesetRequest(featureset_name=feature_set_name)
-            )
-        except Exception as e:
-            raise QwakException(
-                f"Failed to pause Featureset {feature_set_name}, nested exception: [{e.details()}]"
-            ) from e
+        self._features_service.PauseFeatureset(
+            PauseFeaturesetRequest(featureset_name=feature_set_name)
+        )
 
+    @grpc_try_catch_wrapper(
+        "Failed to resume Featureset {feature_set_name}",
+        reraise_non_grpc_error_original_exception=True,
+    )
     def resume_feature_set(self, feature_set_name: str):
-        try:
-            self._features_service.ResumeFeatureset(
-                ResumeFeaturesetRequest(featureset_name=feature_set_name)
+        self._features_service.ResumeFeatureset(
+            ResumeFeaturesetRequest(featureset_name=feature_set_name)
+        )
+
+    @grpc_try_catch_wrapper(
+        "Failed list Featureset versions for Featureset: {featureset_name}",
+        reraise_non_grpc_error_original_exception=True,
+    )
+    def list_featureset_versions(
+        self, featureset_name: str
+    ) -> ListFeaturesetVersionsByNameResponse:
+        return self._features_service.ListFeaturesetVersionsByName(
+            ListFeaturesetVersionsByNameRequest(featureset_name=featureset_name)
+        )
+
+    @grpc_try_catch_wrapper(
+        "Failed to set active version: {version_number} for Featureset: {featureset_name}",
+        reraise_non_grpc_error_original_exception=True,
+    )
+    def set_active_featureset_version(self, featureset_name: str, version_number: int):
+        self._features_service.SetActiveFeaturesetVersion(
+            SetActiveFeaturesetVersionRequest(
+                featureset_name=featureset_name,
+                featureset_version_number=version_number,
             )
-        except Exception as e:
-            raise QwakException(
-                f"Failed to resume Featureset {feature_set_name}, nested exception: [{e.details()}]"
-            ) from e
+        )
+
+    @grpc_try_catch_wrapper(
+        "Failed to delete version: {version_number} for Featureset: {featureset_name}",
+        reraise_non_grpc_error_original_exception=True,
+    )
+    def delete_featureset_version(self, featureset_name: str, version_number: int):
+        self._features_service.DeleteFeaturesetVersion(
+            DeleteFeaturesetVersionRequest(
+                featureset_name=featureset_name,
+                featureset_version_number=version_number,
+            )
+        )
+
+    @grpc_try_catch_wrapper(
+        "Failed to get env to featuresets mapping",
+        reraise_non_grpc_error_original_exception=True,
+    )
+    def get_env_to_featuresets_mapping(self) -> GetEnvToFeatureSetsMappingResponse:
+        return self._features_service.GetEnvToFeatureSetsMapping(
+            GetEnvToFeatureSetsMappingRequest()
+        )