frogml-core 0.0.1__py3-none-any.whl

frogml_services_mock/mocks/ecosystem_service_api.py

@@ -0,0 +1,45 @@
+from _frogml_proto.qwak.ecosystem.v0.credentials_pb2 import (
+    AwsTemporaryCredentials,
+    CloudCredentials,
+)
+from _frogml_proto.qwak.ecosystem.v0.ecosystem_runtime_service_pb2 import (
+    GetAuthenticatedUserContextResponse,
+    GetCloudCredentialsResponse,
+)
+from _frogml_proto.qwak.ecosystem.v0.ecosystem_runtime_service_pb2_grpc import (
+    QwakEcosystemRuntimeServicer,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class EcoSystemServiceMock(QwakEcosystemRuntimeServicer):
+    def __init__(self):
+        self._aws_temp_credentials = None
+        self._authenticated_user_context = None
+        super(EcoSystemServiceMock, self).__init__()
+
+    def given_credentials(self, aws_temp_credentials: AwsTemporaryCredentials):
+        self._aws_temp_credentials: AwsTemporaryCredentials = aws_temp_credentials
+
+    def GetCloudCredentials(self, request, context):
+        """get cloud credentials"""
+        try:
+            return GetCloudCredentialsResponse(
+                cloud_credentials=CloudCredentials(
+                    aws_temporary_credentials=self._aws_temp_credentials
+                )
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def given_authenticated_user_context(
+        self, authenticated_user_context: GetAuthenticatedUserContextResponse
+    ):
+        self._authenticated_user_context = authenticated_user_context
+
+    def GetAuthenticatedUserContext(
+        self, request, context
+    ) -> GetAuthenticatedUserContextResponse:
+        return self._authenticated_user_context

frogml_services_mock/mocks/execution_management_service.py

@@ -0,0 +1,26 @@
+from grpc import RpcError
+
+from _frogml_proto.qwak.execution.v1.execution_service_pb2 import TriggerBackfillResponse
+from _frogml_proto.qwak.execution.v1.execution_service_pb2_grpc import (
+    FeatureStoreExecutionServiceServicer,
+)
+
+
+class ExecutionManagementServiceMock(FeatureStoreExecutionServiceServicer):
+    def __init__(self):
+        self._execution_id = None
+        self._raise_exception_on_request = False
+
+    def set_execution_id(self, execution_id: str) -> None:
+        self._execution_id = execution_id
+
+    def set_raise_exception(self):
+        self._raise_exception_on_request = True
+
+    def clear_raise_exception(self):
+        self._raise_exception_on_request = False
+
+    def TriggerBatchBackfill(self, request, context):
+        if self._raise_exception_on_request:
+            raise RpcError
+        return TriggerBackfillResponse(execution_id=self._execution_id)

frogml_services_mock/mocks/feature_store_data_sources_manager_api.py

@@ -0,0 +1,99 @@
+import time
+import uuid
+
+import grpc
+
+from _frogml_proto.qwak.feature_store.sources.data_source_pb2 import (
+    DataSource,
+    DataSourceDefinition,
+)
+from _frogml_proto.qwak.feature_store.sources.data_source_service_pb2 import (
+    CreateDataSourceResponse,
+    CreateDataSourceUploadURLRequest,
+    CreateDataSourceUploadURLResponse,
+    GetDataSourceByIdRequest,
+    GetDataSourceByIdResponse,
+    GetDataSourceByNameResponse,
+)
+from _frogml_proto.qwak.feature_store.sources.data_source_service_pb2_grpc import (
+    DataSourceServiceServicer,
+)
+
+
+class DataSourceServiceMock(DataSourceServiceServicer):
+    def __init__(self):
+        self._data_sources_spec = {}
+        self._data_source_id_to_name = {}
+
+    def CreateDataSource(self, request, context):
+        data_source_type = request.data_source_spec.WhichOneof("type")
+        data_source_id = str(uuid.uuid4())
+
+        ds_name = (
+            request.data_source_spec.batch_source.name
+            if data_source_type == "batch_source"
+            else request.data_source_spec.stream_source.name
+        )
+
+        self._data_sources_spec[ds_name] = request.data_source_spec
+        self._data_source_id_to_name[data_source_id] = ds_name
+
+        return CreateDataSourceResponse(
+            data_source=DataSource(
+                data_source_definition=DataSourceDefinition(
+                    data_source_id=data_source_id,
+                    data_source_spec=request.data_source_spec,
+                ),
+                metadata=None,
+                feature_sets=[],
+            )
+        )
+
+    def GetDataSourceByName(self, request, context):
+        if request.data_source_name not in self._data_sources_spec:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return GetDataSourceByNameResponse()
+
+        return GetDataSourceByNameResponse(
+            data_source=DataSource(
+                data_source_definition=DataSourceDefinition(
+                    data_source_id="123",
+                    data_source_spec=self._data_sources_spec[request.data_source_name],
+                ),
+                metadata=None,
+                feature_sets=[],
+            )
+        )
+
+    def GetDataSourceById(self, request: GetDataSourceByIdRequest, context):
+        if request.data_source_id not in self._data_source_id_to_name:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return GetDataSourceByIdResponse()
+
+        data_source_name = self._data_source_id_to_name[request.data_source_id]
+        data_source_spec = self._data_sources_spec[data_source_name]
+
+        return GetDataSourceByIdResponse(
+            data_source=DataSource(
+                data_source_definition=DataSourceDefinition(
+                    data_source_id=request.data_source_id,
+                    data_source_spec=data_source_spec,
+                ),
+                metadata=None,
+                feature_sets=[],
+            )
+        )
+
+    def CreateDataSourceUploadURL(
+        self, request: CreateDataSourceUploadURLRequest, context
+    ):
+        return CreateDataSourceUploadURLResponse(
+            upload_url=f"https://data_source_artifacts.s3.amazonaws.com/{request.data_source_name}_{request.object_name}"
+        )
+
+    def CreateDataSourceUploadURLV1(
+        self, request: CreateDataSourceUploadURLRequest, context
+    ):
+        return CreateDataSourceUploadURLResponse(
+            upload_url=f"https://data_source_artifacts.s3.amazonaws.com/{request.data_source_name}_{request.object_name}_{int(time.time() * 1000)}"
+        )

frogml_services_mock/mocks/feature_store_entities_manager_api.py

@@ -0,0 +1,99 @@
+import uuid
+
+import grpc
+
+from _frogml_proto.qwak.feature_store.entities.entity_pb2 import (
+    Entity,
+    EntityDefinition,
+    EntityMetadata,
+    FeatureSetBrief,
+)
+from _frogml_proto.qwak.feature_store.entities.entity_service_pb2 import (
+    CreateEntityResponse,
+    DeleteEntityRequest,
+    DeleteEntityResponse,
+    GetEntityByIdResponse,
+    GetEntityByNameResponse,
+    ListEntitiesResponse,
+)
+from _frogml_proto.qwak.feature_store.entities.entity_service_pb2_grpc import (
+    EntityServiceServicer,
+)
+
+
+class EntityServiceMock(EntityServiceServicer):
+    def __init__(self):
+        self._entities_spec = {}
+        self._entity_id_name = {}
+
+    def reset_entities(self):
+        self._entities_spec.clear()
+
+    def CreateEntity(self, request, context):
+        entity_id = str(uuid.uuid4())
+        self._entities_spec[request.entity_spec.name] = request.entity_spec
+        self._entity_id_name[entity_id] = request.entity_spec.name
+        return CreateEntityResponse()
+
+    def DeleteEntity(self, request: DeleteEntityRequest, context):
+        entity_id: str = request.entity_id
+
+        if entity_id not in self._entity_id_name.keys():
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+        else:
+            entity_name: str = self._entity_id_name[entity_id]
+            del self._entity_id_name[entity_id]
+            del self._entities_spec[entity_name]
+
+            return DeleteEntityResponse()
+
+    def GetEntityByName(self, request, context):
+        if request.entity_name not in self._entities_spec:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return GetEntityByNameResponse()
+        else:
+            entity_key: str = list(self._entity_id_name.keys())[
+                list(self._entity_id_name.values()).index(request.entity_name)
+            ]
+            return GetEntityByNameResponse(
+                entity=Entity(
+                    entity_definition=EntityDefinition(
+                        entity_id=entity_key,
+                        entity_spec=self._entities_spec[request.entity_name],
+                    ),
+                    metadata=None,
+                    feature_sets=[],
+                )
+            )
+
+    def GetEntityById(self, request, context):
+        if request.entity_id not in self._entity_id_name:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return GetEntityByIdResponse()
+        else:
+            entity_name: str = self._entity_id_name[request.entity_id]
+            return GetEntityByIdResponse(
+                entity=Entity(
+                    entity_definition=EntityDefinition(
+                        entity_id=request.entity_id,
+                        entity_spec=self._entities_spec[entity_name],
+                    ),
+                    metadata=None,
+                    feature_sets=[],
+                )
+            )
+
+    def ListEntities(self, request, context):
+        return ListEntitiesResponse(
+            entities=[
+                Entity(
+                    entity_definition=EntityDefinition(
+                        entity_spec=self._entities_spec[entity_name],
+                        entity_id=entity_id,
+                    ),
+                    metadata=EntityMetadata(),
+                    feature_sets=[FeatureSetBrief()],
+                )
+                for entity_id, entity_name in self._entity_id_name.items()
+            ]
+        )

frogml_services_mock/mocks/feature_store_feature_set_manager_api.py

@@ -0,0 +1,269 @@
+import uuid
+from typing import Dict, Optional
+
+import grpc
+
+from _frogml_proto.qwak.feature_store.features.feature_set_pb2 import (
+    DeployedModelInUseLink,
+    FeatureSet,
+    FeatureSetDefinition,
+    FeaturesetInfo,
+    FeatureSetList,
+    FeatureSetMetadata,
+    FeaturesetSchedulingState,
+    FeatureSetSpec,
+    FeatureStatus,
+)
+from _frogml_proto.qwak.feature_store.features.feature_set_service_pb2 import (
+    DeleteFeatureSetResponse,
+    DeleteFeaturesetVersionResponse,
+    GetEnvToFeatureSetsMappingResponse,
+    GetFeatureSetByNameResponse,
+    GetFeaturesetSchedulingStateResponse,
+    ListFeatureSetsResponse,
+    ListFeaturesetVersionsByNameResponse,
+    RegisterFeatureSetResponse,
+    SetActiveFeaturesetVersionResponse,
+    UpdateFeatureSetResponse,
+)
+from _frogml_proto.qwak.feature_store.features.feature_set_service_pb2_grpc import (
+    FeatureSetServiceServicer,
+)
+
+CURRENT_ENV_NAME = "env_name"
+
+
+class FeatureSetServiceMock(FeatureSetServiceServicer):
+    def __init__(self):
+        self._features_spec: Dict[str, FeatureSetSpec] = {}
+        self._fs_next_versions: Dict[str, int] = {}
+        self._fs_versions_by_name: Dict[str, Dict[int, FeatureSetDefinition]] = {}
+        self._fs_definition_by_id: Dict[str, FeatureSetDefinition] = {}
+        # this mapping is used to test the GetEnvToFeatureSetsMapping API. Note that the populated features sets may not be aligned with the self._features_spec featuresets
+        self._env_to_features_spec = (
+            {}
+        )  # env_name -> {featureset_name -> featureset_definition}
+
+    def reset_features(self):
+        self._features_spec.clear()
+        self._fs_next_versions.clear()
+        self._fs_versions_by_name.clear()
+        self._fs_definition_by_id.clear()
+        self._env_to_features_spec.clear()
+
+    def setup_empty_env(self, env_name):
+        self._env_to_features_spec[env_name] = {}
+
+    def register_featureset_to_env_to_features_map(
+        self, env_name, featureset_spec: FeatureSetSpec
+    ):
+        fs_id = str(uuid.uuid4())
+        self._env_to_features_spec[env_name] = self._env_to_features_spec.get(
+            env_name, {}
+        )
+        feature_set_definition = FeatureSetDefinition(
+            feature_set_id=fs_id,
+            feature_set_spec=featureset_spec,
+            status=FeatureStatus.VALID,
+        )
+        self._env_to_features_spec[env_name][
+            featureset_spec.name
+        ] = feature_set_definition
+
+    def RegisterFeatureSet(self, request, context):
+        fs_id = str(uuid.uuid4())
+        curr_version = 1
+        featureset_definition: FeatureSetDefinition = FeatureSetDefinition(
+            feature_set_id=fs_id,
+            feature_set_spec=request.feature_set_spec,
+            featureset_version_number=curr_version,
+            status=FeatureStatus.VALID,
+        )
+        self._features_spec[fs_id] = featureset_definition.feature_set_spec
+        self._fs_definition_by_id[fs_id] = featureset_definition
+        self._fs_versions_by_name[featureset_definition.feature_set_spec.name] = {
+            curr_version: featureset_definition
+        }
+        self._fs_next_versions[featureset_definition.feature_set_spec.name] = (
+            curr_version + 1
+        )
+
+        return RegisterFeatureSetResponse(
+            feature_set=FeatureSet(feature_set_definition=featureset_definition)
+        )
+
+    def UpdateFeatureSet(self, request, context):
+        fs_id: str = request.feature_set_id
+        if fs_id in self._fs_definition_by_id:
+            curr_version = self._fs_next_versions[request.feature_set_spec.name]
+            featureset_definition: FeatureSetDefinition = FeatureSetDefinition(
+                feature_set_id=fs_id,
+                feature_set_spec=request.feature_set_spec,
+                featureset_version_number=curr_version,
+                status=FeatureStatus.VALID,
+            )
+            self._fs_definition_by_id[fs_id] = featureset_definition
+            self._features_spec[fs_id] = featureset_definition.feature_set_spec
+            self._fs_next_versions[featureset_definition.feature_set_spec.name] = (
+                curr_version + 1
+            )
+
+            curr_versions = self._fs_versions_by_name[request.feature_set_spec.name]
+            curr_versions[curr_version] = featureset_definition
+            self._fs_versions_by_name[
+                featureset_definition.feature_set_spec.name
+            ] = curr_versions
+
+            return UpdateFeatureSetResponse(
+                feature_set=FeatureSet(feature_set_definition=featureset_definition)
+            )
+
+        context.set_details(f"Feature set ID {request.feature_set_id} doesn't exist'")
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def DeleteFeatureSet(self, request, context):
+        fs_definition: Optional[FeatureSetDefinition] = self._fs_definition_by_id.get(
+            request.feature_set_id, None
+        )
+        if fs_definition:
+            self._features_spec.pop(request.feature_set_id)
+            self._fs_definition_by_id.pop(request.feature_set_id)
+            self._fs_versions_by_name.pop(fs_definition.feature_set_spec.name)
+            self._fs_next_versions.pop(fs_definition.feature_set_spec.name)
+            return DeleteFeatureSetResponse()
+
+        context.set_details(f"Feature set ID {request.feature_set_id} doesn't exist'")
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def GetFeatureSetByName(self, request, context):
+        feature_sets = [
+            (fs_id, fs_definition)
+            for (fs_id, fs_definition) in self._fs_definition_by_id.items()
+            if fs_definition.feature_set_spec.name == request.feature_set_name
+        ]
+        if feature_sets:
+            fs_id, fs_definition = feature_sets[0]
+            return GetFeatureSetByNameResponse(
+                feature_set=FeatureSet(
+                    feature_set_definition=fs_definition,
+                    metadata=FeatureSetMetadata(),
+                    deployed_models_in_use_link=[DeployedModelInUseLink()],
+                )
+            )
+
+        context.set_details(
+            f"Feature set named {request.feature_set_name} doesn't exist'"
+        )
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def ListFeatureSets(self, request, context):
+        return ListFeatureSetsResponse(
+            feature_families=[
+                FeatureSet(
+                    feature_set_definition=fs_definition,
+                    metadata=FeatureSetMetadata(),
+                    deployed_models_in_use_link=[DeployedModelInUseLink()],
+                )
+                for fs_definition in self._fs_definition_by_id.values()
+            ]
+        )
+
+    def GetFeaturesetSchedulingState(self, request, context):
+        return GetFeaturesetSchedulingStateResponse(
+            state=FeaturesetSchedulingState.SCHEDULING_STATE_ENABLED
+        )
+
+    def GetEnvToFeatureSetsMapping(self, request, context):
+        env_to_featuresets_mapping = {}
+
+        for env_name, featureset_df_dict in self._env_to_features_spec.items():
+            env_to_featuresets_mapping[env_name] = FeatureSetList(
+                feature_sets=[
+                    FeatureSet(
+                        feature_set_definition=fs_definition,
+                        metadata=FeatureSetMetadata(),
+                        deployed_models_in_use_link=[DeployedModelInUseLink()],
+                    )
+                    for fs_name, fs_definition in featureset_df_dict.items()
+                ]
+            )
+
+        return GetEnvToFeatureSetsMappingResponse(
+            env_to_feature_set_mapping=env_to_featuresets_mapping
+        )
+
+    def DeleteFeaturesetVersion(self, request, context):
+        version_number = request.featureset_version_number
+        version_definitions: Optional[
+            Dict[int, FeatureSetDefinition]
+        ] = self._fs_versions_by_name.get(request.featureset_name, None)
+        if version_definitions:
+            version_definition = version_definitions.get(version_number, None)
+            if version_definition:
+                if (
+                    self._fs_definition_by_id[
+                        version_definition.feature_set_id
+                    ].featureset_version_number
+                    != version_number
+                ):
+                    version_definitions.pop(version_number)
+                    return DeleteFeaturesetVersionResponse()
+                else:
+                    context.set_details("Active version cant be deleted")
+                    context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
+            else:
+                context.set_details(
+                    f"Version {version_number} not found for Feature set name {request.featureset_name}"
+                )
+                context.set_code(grpc.StatusCode.NOT_FOUND)
+        else:
+            context.set_details(
+                f"Feature set name {request.featureset_name} doesn't exist'"
+            )
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def SetActiveFeaturesetVersion(self, request, context):
+        version_definitions: Optional[
+            Dict[int, FeatureSetDefinition]
+        ] = self._fs_versions_by_name.get(request.featureset_name, None)
+        if version_definitions:
+            version_definition = version_definitions.get(
+                request.featureset_version_number, None
+            )
+            if version_definition:
+                self._fs_definition_by_id[
+                    version_definition.feature_set_id
+                ] = version_definition
+                return SetActiveFeaturesetVersionResponse()
+            else:
+                context.set_details(
+                    f"Version {request.featureset_version_number} not found for Feature set name {request.featureset_name}"
+                )
+                context.set_code(grpc.StatusCode.NOT_FOUND)
+        else:
+            context.set_details(
+                f"Feature set name {request.featureset_name} doesn't exist'"
+            )
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def ListFeaturesetVersionsByName(self, request, context):
+        versions: Optional[
+            Dict[int, FeatureSetDefinition]
+        ] = self._fs_versions_by_name.get(request.featureset_name, None)
+        if versions:
+            return ListFeaturesetVersionsByNameResponse(
+                featuresets=[
+                    FeaturesetInfo(
+                        featureset_definition=feature_definition,
+                        metadata=FeatureSetMetadata(),
+                    )
+                    for feature_definition in self._fs_versions_by_name[
+                        request.featureset_name
+                    ].values()
+                ]
+            )
+
+        context.set_details(
+            f"Feature set name {request.featureset_name} doesn't exist'"
+        )
+        context.set_code(grpc.StatusCode.NOT_FOUND)