frogml-core 0.0.1__py3-none-any.whl

frogml/feature_store/data_sources/streaming/kafka/authentication.py

@@ -0,0 +1,124 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Type
+
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import (
+    Authentication as ProtoAuthentication,
+)
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import Plain as ProtoPlain
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import Sasl as ProtoSasl
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import (
+    SaslMechanism as ProtoSaslMechanism,
+)
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import (
+    SecurityProtocol as ProtoSecurityProtocol,
+)
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import Ssl as ProtoSsl
+from frogml.clients.secret_service import SecretServiceClient
+from frogml.exceptions import FrogmlException
+
+
+class SaslMechanism(Enum):
+    SCRAMSHA256 = ProtoSaslMechanism.SCRAMSHA256
+    SCRAMSHA512 = ProtoSaslMechanism.SCRAMSHA512
+    PLAIN = ProtoSaslMechanism.PLAIN
+
+
+class SecurityProtocol(Enum):
+    SASL_SSL = ProtoSecurityProtocol.SASL_SSL
+
+
+@dataclass
+class BaseAuthentication(ABC):
+    @abstractmethod
+    def _to_proto(self) -> ProtoAuthentication:
+        pass
+
+    @classmethod
+    def _from_proto(
+        cls,
+        proto_authentication_method: ProtoAuthentication,
+    ) -> Type["BaseAuthentication"]:
+        proto_authentication_method = getattr(
+            proto_authentication_method,
+            proto_authentication_method.WhichOneof("type"),
+        )
+        if isinstance(proto_authentication_method, ProtoPlain):
+            return PlainAuthentication._from_proto(proto_authentication_method)
+        elif isinstance(proto_authentication_method, ProtoSsl):
+            return SslAuthentication._from_proto(proto_authentication_method)
+        elif isinstance(proto_authentication_method, ProtoSasl):
+            return SaslAuthentication._from_proto(proto_authentication_method)
+        else:
+            raise FrogmlException(
+                f"Got unsupported authentication method {proto_authentication_method}"
+            )
+
+
+@dataclass
+class PlainAuthentication(BaseAuthentication):
+    def _to_proto(self) -> ProtoAuthentication:
+        return ProtoAuthentication(plain_configuration=ProtoPlain())
+
+    @classmethod
+    def _from_proto(
+        cls, proto_authentication_method: ProtoPlain
+    ) -> Type["PlainAuthentication"]:
+        return cls()
+
+
+@dataclass
+class SslAuthentication(BaseAuthentication):
+    def _to_proto(self) -> ProtoAuthentication:
+        return ProtoAuthentication(ssl_configuration=ProtoSsl())
+
+    @classmethod
+    def _from_proto(
+        cls, proto_authentication_method: ProtoSsl
+    ) -> Type["SslAuthentication"]:
+        return cls()
+
+
+@dataclass
+class SaslAuthentication(BaseAuthentication):
+    username_secret: str
+    password_secret: str
+    sasl_mechanism: SaslMechanism
+    security_protocol: SecurityProtocol
+
+    def _to_proto(self) -> ProtoAuthentication:
+        self._validate()
+        return ProtoAuthentication(
+            sasl_configuration=ProtoSasl(
+                username_secret=self.username_secret,
+                password_secret=self.password_secret,
+                sasl_mechanism=self.sasl_mechanism.value,
+                security_protocol=self.security_protocol.value,
+            )
+        )
+
+    @classmethod
+    def _from_proto(
+        cls, proto_authentication_method: ProtoSasl
+    ) -> Type["SaslAuthentication"]:
+        return cls(
+            username_secret=proto_authentication_method.username_secret,
+            password_secret=proto_authentication_method.password_secret,
+            sasl_mechanism=SaslMechanism(proto_authentication_method.sasl_mechanism),
+            security_protocol=SecurityProtocol(
+                proto_authentication_method.security_protocol
+            ),
+        )
+
+    def _validate(self):
+        secret_service_client = SecretServiceClient()
+
+        if not secret_service_client.get_secret(self.username_secret):
+            raise FrogmlException(
+                f"Secret for username {self.username_secret} does not exist"
+            )
+        if not secret_service_client.get_secret(self.password_secret):
+            raise FrogmlException(
+                f"Secret for password {self.password_secret} does not exist"
+            )

frogml/feature_store/data_sources/streaming/kafka/deserialization.py

@@ -0,0 +1,115 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Callable, Optional, Type
+
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import (
+    CustomDeserializer as ProtoCustomDeserializer,
+)
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import (
+    Deserialization as ProtoDeserialization,
+)
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import (
+    GenericDeserializer as ProtoGenericDeserializer,
+)
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import (
+    MessageFormat as ProtoMessageFormat,
+)
+from frogml.exceptions import FrogmlException
+
+
+class MessageFormat(Enum):
+    JSON = ProtoMessageFormat.JSON
+    AVRO = ProtoMessageFormat.AVRO
+
+
+class Deserializer(ABC):
+    @abstractmethod
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoDeserialization:
+        pass
+
+    @abstractmethod
+    def _get_function(self) -> Optional[Callable]:
+        pass
+
+    @classmethod
+    def _from_proto(
+        cls,
+        proto_deserializer: ProtoDeserialization,
+    ) -> Type["Deserializer"]:
+        deserializer = getattr(
+            proto_deserializer, proto_deserializer.WhichOneof("type")
+        )
+
+        if isinstance(deserializer, ProtoGenericDeserializer):
+            return GenericDeserializer._from_proto(proto_deserializer=deserializer)
+        elif isinstance(deserializer, ProtoCustomDeserializer):
+            return CustomDeserializer._from_proto(proto_deserializer=deserializer)
+        else:
+            raise FrogmlException(f"Got unsupported deserializer type {deserializer}")
+
+
+@dataclass
+class GenericDeserializer(Deserializer):
+    message_format: MessageFormat
+    schema: str
+
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoDeserialization:
+        # TODO: add backend schema validation
+        return ProtoDeserialization(
+            generic_deserializer=ProtoGenericDeserializer(
+                deserializer_format=self.message_format.value, schema=self.schema
+            )
+        )
+
+    def _get_function(self) -> Optional[Callable]:
+        return None
+
+    @classmethod
+    def _from_proto(
+        cls, proto_deserializer: ProtoGenericDeserializer
+    ) -> Type["GenericDeserializer"]:
+        return cls(
+            message_format=MessageFormat(proto_deserializer.deserializer_format),
+            schema=proto_deserializer.schema,
+        )
+
+
+@dataclass
+class CustomDeserializer(Deserializer):
+    function: Callable
+    _artifact_path: Optional[str] = field(init=False, default=None)
+
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoDeserialization:
+        self._validate()
+
+        if artifact_path:
+            self._artifact_path = artifact_path
+
+        # TODO: add backend schema validation
+        return ProtoDeserialization(
+            custom_deserializer=ProtoCustomDeserializer(
+                function_name=self.function.__name__, artifact_path=self._artifact_path
+            )
+        )
+
+    @classmethod
+    def _from_proto(
+        cls, proto_deserializer: ProtoCustomDeserializer
+    ) -> Type["CustomDeserializer"]:
+        def dummy_deserializer(df):
+            return df
+
+        custom_function: Callable = dummy_deserializer
+        custom_function.__name__ = proto_deserializer.function_name
+
+        return cls(function=custom_function)
+
+    def _validate(self):
+        if self.function is None:
+            raise FrogmlException("deserialization function must be set!")
+        if self.function.__name__ == "<lambda>":
+            raise FrogmlException("Custom Deserializer can not be set with a lambda")
+
+    def _get_function(self) -> Optional[Callable]:
+        return self.function

frogml/feature_store/data_sources/streaming/kafka/kafka.py

@@ -0,0 +1,121 @@
+import inspect
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Dict, Optional, Type
+
+from _frogml_proto.qwak.feature_store.sources.data_source_pb2 import (
+    DataSourceSpec as ProtoDataSourceSpec,
+)
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import (
+    KafkaSourceV1 as ProtoKafkaSourceV1,
+)
+from _frogml_proto.qwak.feature_store.sources.streaming_pb2 import (
+    StreamingSource as ProtoStreamingSource,
+)
+from frogml.exceptions import FrogmlException
+from frogml.feature_store._common.artifact_utils import ArtifactSpec
+from frogml.feature_store.data_sources.streaming._streaming import BaseStreamingSource
+from frogml.feature_store.data_sources.streaming.kafka.authentication import (
+    BaseAuthentication,
+    SslAuthentication,
+)
+from frogml.feature_store.data_sources.streaming.kafka.deserialization import (
+    Deserializer,
+)
+
+
+@dataclass
+class KafkaSource(BaseStreamingSource):
+    bootstrap_servers: str
+
+    # Deserialization
+    deserialization: Deserializer
+
+    # secret configs, the value is resolved to the secret,
+    # s.t. (key, value) -> (key, get_secret(value))
+    # not all configs will be respected, this is a best-effort
+    secret_configs: Dict[str, str] = field(default_factory=lambda: {})
+
+    # passthrough configs - not all configs will be respected,
+    # this is a best-effort
+    passthrough_configs: Dict[str, str] = field(default_factory=lambda: {})
+
+    # the following 3 are pairwise mutually exclusive
+    assign: Optional[str] = None
+    subscribe: Optional[str] = None
+    subscribe_pattern: Optional[str] = None
+
+    authentication_method: Optional[BaseAuthentication] = field(
+        default_factory=lambda: SslAuthentication()
+    )
+
+    def __post_init__(self):
+        self._validate()
+
+    def _validate(self):
+        num_defined = len(
+            [
+                _
+                for _ in [self.assign, self.subscribe, self.subscribe_pattern]
+                if _ is not None
+            ]
+        )
+        if num_defined != 1:
+            raise FrogmlException(
+                "Exactly one of (assign, subscribe, subscribe_pattern) must be defined!"
+            )
+
+    def _get_artifacts(self) -> Optional["ArtifactSpec"]:
+        deserializer_function: Optional[Callable] = self.deserialization._get_function()
+
+        if deserializer_function:
+            return ArtifactSpec(
+                artifact_name=self.name,
+                root_module_path=Path(
+                    os.path.abspath(inspect.getfile(deserializer_function))
+                ).parent,
+                artifact_object=self.deserialization,
+                callables=[deserializer_function],
+                suffix="resource.zip",
+            )
+
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
+        return ProtoDataSourceSpec(
+            stream_source=ProtoStreamingSource(
+                name=self.name,
+                description=self.description,
+                kafkaSourceV1=ProtoKafkaSourceV1(
+                    bootstrap_servers=self.bootstrap_servers,
+                    assign=self.assign,
+                    subscribe=self.subscribe,
+                    subscribe_pattern=self.subscribe_pattern,
+                    secret_configs=self.secret_configs,
+                    passthrough_configs=self.passthrough_configs,
+                    authentication_method=self.authentication_method._to_proto(),
+                    deserialization=self.deserialization._to_proto(
+                        artifact_path=artifact_url
+                    ),
+                ),
+            )
+        )
+
+    @classmethod
+    def _from_proto(cls, proto) -> Type["KafkaSource"]:
+        kafka = proto.kafkaSourceV1
+        topic_configuration_key = kafka.WhichOneof("topic_configuration")
+        oneof_args = {topic_configuration_key: getattr(kafka, topic_configuration_key)}
+        return cls(
+            name=proto.name,
+            description=proto.description,
+            bootstrap_servers=kafka.bootstrap_servers,
+            secret_configs=kafka.secret_configs,
+            passthrough_configs=kafka.passthrough_configs,
+            deserialization=Deserializer._from_proto(
+                proto_deserializer=kafka.deserialization
+            ),
+            authentication_method=BaseAuthentication._from_proto(
+                proto_authentication_method=kafka.authentication_method
+            ),
+            **oneof_args,
+        )

frogml/feature_store/data_sources/time_partition_columns.py

@@ -0,0 +1,178 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Optional
+
+from _frogml_proto.qwak.feature_store.sources.batch_pb2 import (
+    DatePartitionColumns as ProtoDatePartitionColumns,
+)
+from _frogml_proto.qwak.feature_store.sources.batch_pb2 import (
+    DayFragmentColumn as ProtoDayFragmentColumn,
+)
+from _frogml_proto.qwak.feature_store.sources.batch_pb2 import (
+    MonthFragmentColumn as ProtoMonthFragmentColumn,
+)
+from _frogml_proto.qwak.feature_store.sources.batch_pb2 import (
+    NumericColumnRepresentation as ProtoNumericColumnRepresentation,
+)
+from _frogml_proto.qwak.feature_store.sources.batch_pb2 import (
+    TextualColumnRepresentation as ProtoTextualColumnRepresentation,
+)
+from _frogml_proto.qwak.feature_store.sources.batch_pb2 import (
+    TimeFragmentedPartitionColumns as ProtoTimeFragmentedPartitionColumns,
+)
+from _frogml_proto.qwak.feature_store.sources.batch_pb2 import (
+    YearFragmentColumn as ProtoYearFragmentColumn,
+)
+
+
+class ColumnRepresentation(Enum):
+    """
+    Time fragment columns representations
+        NumericColumnRepresentation:
+            Year:   2020 | '2020'
+            Month:  2  | '02' | '2'
+            Day:    2  | '02' | '2'
+        TextualColumnRepresentation:
+            Month: 'Jan' | 'JANUARY' (case-insensitive)
+    """
+
+    NumericColumnRepresentation = 1
+    TextualColumnRepresentation = 2
+
+    def _to_proto(self):
+        if self == ColumnRepresentation.NumericColumnRepresentation:
+            res = ProtoNumericColumnRepresentation()
+        elif self == ColumnRepresentation.TextualColumnRepresentation:
+            res = ProtoTextualColumnRepresentation()
+        else:
+            raise ValueError(
+                f"Unsupported ColumnRepresentation: {self.name}, supported types: {ColumnRepresentation.NumericColumnRepresentation}, {ColumnRepresentation.TextualColumnRepresentation} "
+            )
+        return res
+
+
+class FragmentColumn(ABC):
+    column_name: str
+    representation: ColumnRepresentation
+
+    def __init__(self, column_name: str, representation: ColumnRepresentation):
+        if representation not in self._valid_representations:
+            raise ValueError(
+                f"{self.__class__.__name__} doesn't support representation: {representation}, supported types: {self._valid_representations} "
+            )
+        self.column_name = column_name
+        self.representation = representation
+
+    @property
+    @abstractmethod
+    def _valid_representations(self) -> List[ColumnRepresentation]:
+        pass
+
+    @abstractmethod
+    def _to_proto(self):
+        pass
+
+
+class YearFragmentColumn(FragmentColumn):
+    _valid_representations: List[ColumnRepresentation] = [
+        ColumnRepresentation.NumericColumnRepresentation,
+    ]
+
+    def __init__(self, column_name: str, representation: ColumnRepresentation):
+        super().__init__(column_name=column_name, representation=representation)
+
+    def _to_proto(self):
+        return ProtoYearFragmentColumn(
+            column_name=self.column_name,
+            numeric_column_representation=self.representation._to_proto(),
+        )
+
+
+class MonthFragmentColumn(FragmentColumn):
+    _valid_representations: List[ColumnRepresentation] = [
+        ColumnRepresentation.NumericColumnRepresentation,
+        ColumnRepresentation.TextualColumnRepresentation,
+    ]
+
+    def __init__(self, column_name: str, representation: ColumnRepresentation):
+        super().__init__(column_name=column_name, representation=representation)
+
+    def _to_proto(self):
+        proto_numeric_column_representation = None
+        proto_textual_column_representation = None
+        if self.representation == ColumnRepresentation.NumericColumnRepresentation:
+            proto_numeric_column_representation = self.representation._to_proto()
+        elif self.representation == ColumnRepresentation.TextualColumnRepresentation:
+            proto_textual_column_representation = self.representation._to_proto()
+        else:
+            raise ValueError(
+                f"{self.__class__.__name__} partition doesn't support representation: {self.representation}, supported types: {self._valid_representations} "
+            )
+        return ProtoMonthFragmentColumn(
+            column_name=self.column_name,
+            numeric_column_representation=proto_numeric_column_representation,
+            textual_column_representation=proto_textual_column_representation,
+        )
+
+
+class DayFragmentColumn(FragmentColumn):
+    _valid_representations: List[ColumnRepresentation] = [
+        ColumnRepresentation.NumericColumnRepresentation
+    ]
+
+    def __init__(self, column_name: str, representation: ColumnRepresentation):
+        super().__init__(column_name=column_name, representation=representation)
+
+    def _to_proto(self):
+        return ProtoDayFragmentColumn(
+            column_name=self.column_name,
+            numeric_column_representation=self.representation._to_proto(),
+        )
+
+
+@dataclass
+class TimePartitionColumns(ABC):
+    @abstractmethod
+    def _to_proto(self):
+        pass
+
+
+@dataclass
+class DatePartitionColumns(TimePartitionColumns):
+    date_column_name: str
+    date_format: str
+
+    def _to_proto(self):
+        return ProtoDatePartitionColumns(
+            date_column_name=self.date_column_name,
+            date_format=self.date_format,
+        )
+
+
+@dataclass
+class TimeFragmentedPartitionColumns(TimePartitionColumns):
+    year_partition_column: YearFragmentColumn
+    month_partition_column: Optional[MonthFragmentColumn] = None
+    day_partition_column: Optional[DayFragmentColumn] = None
+
+    def __post_init__(self):
+        if not self.month_partition_column and self.day_partition_column:
+            raise ValueError(
+                "If day partition column is set then month partition column must be set as well"
+            )
+
+    def _to_proto(self):
+        return ProtoTimeFragmentedPartitionColumns(
+            year_partition_column=self.year_partition_column._to_proto(),
+            month_partition_column=(
+                self.month_partition_column._to_proto()
+                if self.month_partition_column
+                else None
+            ),
+            day_partition_column=(
+                self.day_partition_column._to_proto()
+                if self.day_partition_column
+                else None
+            ),
+        )

frogml/feature_store/entities/entity.py

@@ -0,0 +1,73 @@
+from typing import Optional
+
+from _frogml_proto.qwak.feature_store.entities.entity_pb2 import EntitySpec
+from frogml.clients.feature_store import FeatureRegistryClient
+
+
+class Entity:
+    """
+    An entity is an object which features sets can be associated with.
+
+    Features created via the feature store are tied to a specific business entity.
+    For example registered users, transactions, or merchants.
+
+    Example of an Entity definition:
+
+    .. code-block:: python
+
+        from frogml.feature_store import Entity
+
+        customer = Entity(
+            name='organization_id',
+            description='An organization which is a customer of our platform',
+        )
+    """
+
+    def __init__(self, name: str, description: str, key: Optional[str] = None):
+        """
+        Create a new Entity
+
+        :param name: The name of the entity, must be unique
+        :param description: Short human-readable description of the entity
+        :param key: The actual key this entity is referring to. Used when user provides a key instead of
+                    entity configuration. Currently only a single key is supported.
+        """
+        self.name = name
+        self.description = description
+        self.key: Optional[str] = key
+
+    def _to_proto(self):
+        return EntitySpec(
+            name=self.name,
+            keys=[self.name if not self.key else self.key],
+            description=self.description,
+            value_type=1,  # string
+        )
+
+    @classmethod
+    def _from_proto(cls, proto):
+        entity_spec: EntitySpec = proto.entity_spec
+
+        # entity spec keys is a repeated str but only single key is supported
+        keys_repr: str = ",".join(entity_spec.keys)
+
+        return cls(
+            name=entity_spec.name,
+            description=entity_spec.description,
+            key=keys_repr,
+        )
+
+    def register(self):
+        """
+        Explicitly register this entity to Frogml's Feature Store
+        """
+
+        registry = FeatureRegistryClient()
+        existing_entity = registry.get_entity_by_name(self.name)
+        if existing_entity:
+            registry.update_entity(
+                existing_entity.entity.entity_definition.entity_id,
+                self._to_proto(),
+            )
+        else:
+            registry.create_entity(self._to_proto())