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

qwak/feature_store/data_sources/{batch_sources → batch}/vertica.py

@@ -1,15 +1,14 @@
 from dataclasses import dataclass
+from typing import Optional
 
 from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
-)
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     VerticaSource as ProtoVerticaSource,
 )
 from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
     DataSourceSpec as ProtoDataSourceSpec,
 )
-from qwak.feature_store.data_sources.batch_sources._batch import BaseBatchSource
+from qwak.feature_store.data_sources.batch._batch import BaseBatchSource
 
 
 @dataclass
@@ -21,6 +20,7 @@ class VerticaSource(BaseBatchSource):
     table: str
     username_secret_name: str
     password_secret_name: str
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -38,8 +38,9 @@ class VerticaSource(BaseBatchSource):
             table=vertica.table,
         )
 
-    def _to_proto(self):
+    def _to_proto(self, artifact_url: Optional[str] = None):
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -53,5 +54,5 @@ class VerticaSource(BaseBatchSource):
                     port=self.port,
                     table=self.table,
                 ),
-            )
+            ),
         )

qwak/feature_store/data_sources/source_authentication.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
+    AwsAssumeRoleAuthentication as ProtoAwsAssumeRoleAuthentication,
+    AwsCredentialsAuthentication as ProtoAwsCredentialsAuthentication,
+)
+
+
+class AwsAuthentication(ABC):
+    @abstractmethod
+    def _to_proto(self):
+        pass
+
+
+@dataclass
+class AwsAssumeRoleAuthentication(AwsAuthentication):
+    role_arn: str
+
+    def _to_proto(self):
+        return ProtoAwsAssumeRoleAuthentication(role_arn=self.role_arn)
+
+
+@dataclass
+class AwsCredentialsAuthentication(AwsAuthentication):
+    access_key_secret_name: str
+    secret_key_secret_name: str
+
+    def _to_proto(self):
+        return ProtoAwsCredentialsAuthentication(
+            access_key_secret_name=self.access_key_secret_name,
+            secret_key_secret_name=self.secret_key_secret_name,
+        )

qwak/feature_store/data_sources/streaming/_streaming.py

@@ -0,0 +1,9 @@
+from abc import ABC
+from dataclasses import dataclass
+
+from qwak.feature_store.data_sources.base import BaseSource
+
+
+@dataclass
+class BaseStreamingSource(BaseSource, ABC):
+    pass

qwak/feature_store/data_sources/streaming/kafka/__init__.py

@@ -0,0 +1,25 @@
+from qwak.feature_store.data_sources.streaming.kafka.authentication import (
+    PlainAuthentication,
+    SaslAuthentication,
+    SaslMechanism,
+    SecurityProtocol,
+    SslAuthentication,
+)
+from qwak.feature_store.data_sources.streaming.kafka.deserialization import (
+    CustomDeserializer,
+    GenericDeserializer,
+    MessageFormat,
+)
+from qwak.feature_store.data_sources.streaming.kafka.kafka import KafkaSource
+
+__all__ = [
+    "KafkaSource",
+    "CustomDeserializer",
+    "GenericDeserializer",
+    "SaslAuthentication",
+    "SslAuthentication",
+    "PlainAuthentication",
+    "MessageFormat",
+    "SaslMechanism",
+    "SecurityProtocol",
+]

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

@@ -0,0 +1,120 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Type
+
+from _qwak_proto.qwak.feature_store.sources.streaming_pb2 import (
+    Authentication as ProtoAuthentication,
+    Plain as ProtoPlain,
+    Sasl as ProtoSasl,
+    SaslMechanism as ProtoSaslMechanism,
+    SecurityProtocol as ProtoSecurityProtocol,
+    Ssl as ProtoSsl,
+)
+from qwak.clients.secret_service import SecretServiceClient
+from qwak.exceptions import QwakException
+
+
+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 QwakException(
+                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 QwakException(
+                f"Secret for username {self.username_secret} does not exist"
+            )
+        if not secret_service_client.get_secret(self.password_secret):
+            raise QwakException(
+                f"Secret for password {self.password_secret} does not exist"
+            )

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

@@ -0,0 +1,109 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Callable, Optional, Type
+
+from _qwak_proto.qwak.feature_store.sources.streaming_pb2 import (
+    CustomDeserializer as ProtoCustomDeserializer,
+    Deserialization as ProtoDeserialization,
+    GenericDeserializer as ProtoGenericDeserializer,
+    MessageFormat as ProtoMessageFormat,
+)
+from qwak.exceptions import QwakException
+
+
+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 QwakException(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 QwakException("deserialization function must be set!")
+        if self.function.__name__ == "<lambda>":
+            raise QwakException("Custom Deserializer can not be set with a lambda")
+
+    def _get_function(self) -> Optional[Callable]:
+        return self.function

qwak/feature_store/data_sources/streaming/kafka/kafka.py

@@ -0,0 +1,120 @@
+import inspect
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Dict, Optional, Type
+
+from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
+    DataSourceSpec as ProtoDataSourceSpec,
+)
+from _qwak_proto.qwak.feature_store.sources.streaming_pb2 import (
+    KafkaSourceV1 as ProtoKafkaSourceV1,
+    StreamingSource as ProtoStreamingSource,
+)
+from qwak.exceptions import QwakException
+from qwak.feature_store._common.artifact_utils import ArtifactSpec
+from qwak.feature_store.data_sources.streaming._streaming import BaseStreamingSource
+from qwak.feature_store.data_sources.streaming.kafka.authentication import (
+    BaseAuthentication,
+    SslAuthentication,
+)
+from qwak.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()
+    )
+
+    repository: Optional[str] = None
+
+    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 QwakException(
+                "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(
+            data_source_repository_name=self.repository,
+            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,
+        )

qwak/feature_store/data_sources/time_partition_columns.py

@@ -0,0 +1,166 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Optional
+
+from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
+    DatePartitionColumns as ProtoDatePartitionColumns,
+    DayFragmentColumn as ProtoDayFragmentColumn,
+    MonthFragmentColumn as ProtoMonthFragmentColumn,
+    NumericColumnRepresentation as ProtoNumericColumnRepresentation,
+    TextualColumnRepresentation as ProtoTextualColumnRepresentation,
+    TimeFragmentedPartitionColumns as ProtoTimeFragmentedPartitionColumns,
+    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
+            ),
+        )

qwak/feature_store/entities/entity.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from _qwak_proto.qwak.feature_store.entities.entity_pb2 import EntitySpec
 from qwak.clients.feature_store import FeatureRegistryClient
 
@@ -21,28 +23,39 @@ class Entity:
         )
     """
 
-    def __init__(self, name: str, description: str):
+    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],
+            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 = proto.entity_spec
-        return cls(name=entity_spec.name, description=entity_spec.description)
+        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):
         """