frogml-core 0.0.116__py3-none-any.whl → 0.0.117__py3-none-any.whl

frogml_services_mock/mocks/ecosystem_service_api.py

@@ -1,45 +0,0 @@
-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 as FrogmlEcosystemRuntimeServicer,
-)
-from frogml_services_mock.mocks.utils.exception_handlers import (
-    raise_internal_grpc_error,
-)
-
-
-class EcoSystemServiceMock(FrogmlEcosystemRuntimeServicer):
-    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

@@ -1,26 +0,0 @@
-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

@@ -1,99 +0,0 @@
-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

@@ -1,99 +0,0 @@
-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

@@ -1,269 +0,0 @@
-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)