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

qwak/feature_store/data_sources/base.py

@@ -0,0 +1,125 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional, Tuple
+
+from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
+    DataSourceSpec as ProtoDataSourceSpec,
+)
+from _qwak_proto.qwak.feature_store.v1.internal.data_source.data_source_service_pb2 import (
+    GetDataSourceSourceCodeUploadResponse,
+)
+from qwak.clients.feature_store import FeatureRegistryClient
+from qwak.exceptions import QwakException
+from qwak.feature_store._common.artifact_utils import ArtifactSpec, ArtifactsUploader
+from qwak.feature_store._common.source_code_spec import SourceCodeSpec
+from qwak.feature_store._common.source_code_spec_factory import SourceCodeSpecFactory
+from qwak.feature_store.data_sources.attributes import DataSourceAttributes
+from qwak.feature_store.validations.validation_options import (
+    DataSourceValidationOptions,
+)
+from qwak.feature_store.validations.validation_response import SuccessValidationResponse
+
+if TYPE_CHECKING:
+    try:
+        import pandas as pd
+    except ImportError:
+        pass
+
+
+@dataclass
+class BaseSource(ABC):
+    name: str
+    description: str
+    attributes: DataSourceAttributes = field(
+        init=False, default_factory=lambda: DataSourceAttributes()
+    )
+
+    @abstractmethod
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
+        """
+        Internal qwak function, converts this BaseSource instance to its protobuf representation.
+        Args:
+            artifact_url: provided if this base source artifact has already been uploaded
+            data_source_definition_path: path to the module instantiating this base source
+        """
+        pass
+
+    def _get_artifacts(self) -> Optional["ArtifactSpec"]:
+        return None
+
+    def _upload_artifact(self) -> Optional[str]:
+        artifact: Optional["ArtifactSpec"] = self._get_artifacts()
+        if artifact:
+            return ArtifactsUploader.upload(artifact)
+
+    def _prepare_and_get(
+        self,
+        artifact_url: Optional[str] = None,
+        source_definition_path: Optional[Path] = None,
+    ) -> Tuple[ProtoDataSourceSpec, Optional[str]]:
+        uploaded_artifact_url: Optional[str] = artifact_url
+
+        if not artifact_url:
+            uploaded_artifact_url = self._upload_artifact()
+
+        if source_definition_path:
+            code_upload_args: (
+                GetDataSourceSourceCodeUploadResponse
+            ) = FeatureRegistryClient().get_datasource_source_code_presign_url(
+                ds_name=self.name, repository_name=getattr(self, "repository")
+            )
+            source_code_spec: SourceCodeSpec = (
+                SourceCodeSpecFactory.get_zip_source_code_spec(
+                    main_entity_path=source_definition_path,
+                    url=code_upload_args.upload_url,
+                    extra_headers=code_upload_args.extra_headers,
+                )
+            )
+
+            self.attributes = DataSourceAttributes(source_code_spec=source_code_spec)
+
+        proto_spec: ProtoDataSourceSpec = self._to_proto(
+            artifact_url=uploaded_artifact_url
+        )
+
+        proto_spec.MergeFrom(
+            ProtoDataSourceSpec(data_source_attributes=self.attributes._to_proto())
+        )
+
+        return proto_spec, uploaded_artifact_url
+
+    @classmethod
+    @abstractmethod
+    def _from_proto(cls, proto):
+        pass
+
+    def get_sample(
+        self,
+        number_of_rows: int = 10,
+        validation_options: Optional[DataSourceValidationOptions] = None,
+    ) -> "pd.DataFrame":
+        """
+        Tries to get a sample of length `number_rows` from the data source.
+        Args:
+            number_of_rows: number of rows to get from data source
+            validation_options: validation options
+        Returns:
+            A tuple containing the resulting dataframe and a tuple of the columns names and types.
+            (the types are pyspark dataframe types)
+        """
+        from qwak.feature_store.validations.validator import FeaturesOperatorValidator
+
+        v = FeaturesOperatorValidator()
+
+        response, _ = v.validate_data_source(
+            data_source=self,
+            sample_size=number_of_rows,
+            validation_options=validation_options,
+            silence_specific_exceptions=False,
+        )
+
+        if isinstance(response, SuccessValidationResponse):
+            return response.sample
+        else:
+            raise QwakException(f"Sampling failed: \n{response}")

qwak/feature_store/data_sources/batch/_batch.py

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

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

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from typing import Optional
 
 from qwak.exceptions import QwakException
-from qwak.feature_store.data_sources.batch_sources._batch import BaseBatchSource
+from qwak.feature_store.data_sources.batch._batch import BaseBatchSource
 
 
 @dataclass
@@ -13,6 +13,7 @@ class JdbcSource(BaseBatchSource, ABC):
     url: Optional[str] = None
     db_table: Optional[str] = None
     query: Optional[str] = None
+    repository: Optional[str] = None
 
     def __post_init__(self):
         self._validate()

qwak/feature_store/data_sources/batch/athena.py

@@ -0,0 +1,260 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
+    AthenaSourceV1 as ProtoAthenaSourceV1,
+    AwsAssumeRoleAuthentication as ProtoAwsAssumeRoleAuthentication,
+    AwsAuthentication as ProtoAwsAuthentication,
+    AwsCredentialsAuthentication as ProtoAwsCredentialsAuthentication,
+    BatchSource as ProtoBatchSource,
+    DatePartitionColumns as ProtoDatePartitionColumns,
+    JdbcSource as ProtoJdbcSource,
+    TimeFragmentedPartitionColumns as ProtoTimeFragmentedPartitionColumns,
+)
+from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
+    DataSourceSpec as ProtoDataSourceSpec,
+)
+from qwak.feature_store.data_sources.batch._jdbc import JdbcSource
+from qwak.feature_store.data_sources.source_authentication import (
+    AwsAssumeRoleAuthentication,
+    AwsAuthentication,
+    AwsCredentialsAuthentication,
+)
+from qwak.feature_store.data_sources.time_partition_columns import (
+    ColumnRepresentation,
+    DatePartitionColumns,
+    DayFragmentColumn,
+    MonthFragmentColumn,
+    TimeFragmentedPartitionColumns,
+    TimePartitionColumns,
+    YearFragmentColumn,
+)
+
+
+@dataclass
+class AthenaSource(JdbcSource):
+    aws_region: str = str()
+    s3_output_location: str = str()
+    aws_authentication: Optional[AwsAuthentication] = None
+    time_partition_columns: Optional[TimePartitionColumns] = None
+    workgroup: Optional[str] = None
+    repository: Optional[str] = None
+
+    def _validate(self):
+        super()._validate()
+
+        if self.aws_authentication is None:
+            raise AttributeError("`aws_authentication` must be set")
+
+        if not self.aws_region:
+            raise AttributeError("`aws_region` must be set")
+
+        if not self.s3_output_location and not self.workgroup:
+            raise AttributeError(
+                "At least one of: `s3_output_location`, `workgroup` must be set"
+            )
+
+    def _to_proto(self, artifact_url: Optional[str] = None):
+        proto_date_partition_columns = None
+        proto_time_fragment_partition_columns = None
+
+        if self.time_partition_columns:
+            proto_partition = self.time_partition_columns._to_proto()
+            if isinstance(proto_partition, ProtoDatePartitionColumns):
+                proto_date_partition_columns = proto_partition
+            elif isinstance(proto_partition, ProtoTimeFragmentedPartitionColumns):
+                proto_time_fragment_partition_columns = proto_partition
+            else:
+                raise ValueError(
+                    f"Got an unsupported time partition column: {proto_partition}"
+                )
+
+        proto_aws_assume_role_authentication = None
+        proto_aws_static_credentials_authentication = None
+
+        proto_aws_authentication = self.aws_authentication._to_proto()
+        if isinstance(proto_aws_authentication, ProtoAwsAssumeRoleAuthentication):
+            proto_aws_assume_role_authentication = proto_aws_authentication
+        elif isinstance(proto_aws_authentication, ProtoAwsCredentialsAuthentication):
+            proto_aws_static_credentials_authentication = proto_aws_authentication
+
+        proto_auth: AwsAuthentication = ProtoAwsAuthentication(
+            aws_assume_role_authentication=proto_aws_assume_role_authentication,
+            aws_static_credentials_authentication=proto_aws_static_credentials_authentication,
+        )
+
+        return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
+            batch_source=ProtoBatchSource(
+                name=self.name,
+                description=self.description,
+                date_created_column=self.date_created_column,
+                jdbcSource=ProtoJdbcSource(
+                    url=self.url,
+                    username_secret_name=self.username_secret_name,
+                    password_secret_name=self.password_secret_name,
+                    db_table=self.db_table,
+                    query=self.query,
+                    athenaSource=ProtoAthenaSourceV1(
+                        aws_authentication=proto_auth,
+                        date_partition_columns=proto_date_partition_columns,
+                        time_fragmented_partition_columns=proto_time_fragment_partition_columns,
+                        s3_output_location=self.s3_output_location,
+                        aws_region=self.aws_region,
+                        workgroup=self.workgroup,
+                    ),
+                ),
+            ),
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoBatchSource):
+        proto_jdbc_source: ProtoJdbcSource = proto.jdbcSource
+        proto_athena_source: ProtoAthenaSourceV1 = proto_jdbc_source.athenaSource
+
+        aws_authentication: AwsAuthentication = (
+            AthenaSource._extract_and_configure_aws_authentication(
+                proto_athena_source.aws_authentication
+            )
+        )
+
+        time_partition_columns: Optional[TimePartitionColumns] = (
+            AthenaSource._extract_partition_column(proto_athena_source)
+        )
+        workgroup: Optional[str] = (
+            proto_athena_source.workgroup
+            if proto_athena_source.HasField("workgroup")
+            else None
+        )
+
+        return cls(
+            name=proto.name,
+            date_created_column=proto.date_created_column,
+            description=proto.description,
+            password_secret_name=proto_jdbc_source.password_secret_name,
+            username_secret_name=proto_jdbc_source.username_secret_name,
+            db_table=proto_jdbc_source.db_table,
+            query=proto_jdbc_source.query,
+            aws_authentication=aws_authentication,
+            time_partition_columns=time_partition_columns,
+            aws_region=proto_athena_source.aws_region,
+            s3_output_location=proto_athena_source.s3_output_location,
+            workgroup=workgroup,
+        )
+
+    @staticmethod
+    def _extract_and_configure_aws_authentication(
+        proto_aws_authentication: ProtoAwsAuthentication,
+    ) -> AwsAuthentication:
+        auth_type = proto_aws_authentication.WhichOneof("type")
+
+        if auth_type == "aws_static_credentials_authentication":
+            proto_aws_authentication = (
+                AthenaSource._extract_and_configure_aws_credentials_auth(
+                    proto_aws_authentication.aws_static_credentials_authentication
+                )
+            )
+        elif auth_type == "aws_assume_role_authentication":
+            proto_aws_authentication = AthenaSource._extract_aws_assume_role_auth(
+                proto_aws_authentication.aws_assume_role_authentication
+            )
+        else:
+            raise ValueError(
+                f"Got an unsupported aws authentication method: {proto_aws_authentication}"
+            )
+
+        return proto_aws_authentication
+
+    @staticmethod
+    def _extract_and_configure_aws_credentials_auth(
+        proto_aws_creds: ProtoAwsCredentialsAuthentication,
+    ) -> AwsCredentialsAuthentication:
+        return AwsCredentialsAuthentication(
+            access_key_secret_name=proto_aws_creds.access_key_secret_name,
+            secret_key_secret_name=proto_aws_creds.secret_key_secret_name,
+        )
+
+    @staticmethod
+    def _extract_aws_assume_role_auth(
+        proto_aws_assume_role: ProtoAwsAssumeRoleAuthentication,
+    ) -> AwsAssumeRoleAuthentication:
+        return AwsAssumeRoleAuthentication(
+            role_arn=proto_aws_assume_role.role_arn,
+        )
+
+    @staticmethod
+    def _extract_partition_column(
+        proto_athena_source: ProtoAthenaSourceV1,
+    ) -> Optional[TimePartitionColumns]:
+        if proto_athena_source.HasField("time_partition_columns"):
+            partition_cols_type = proto_athena_source.WhichOneof(
+                "time_partition_columns"
+            )
+
+            if partition_cols_type == "date_partition_columns":
+                res = DatePartitionColumns(
+                    date_column_name=proto_athena_source.date_partition_columns.date_column_name,
+                    date_format=proto_athena_source.date_partition_columns.date_format,
+                )
+            elif partition_cols_type == "time_fragmented_partition_columns":
+                year_rep = AthenaSource._extract_column_representation(
+                    proto_athena_source.time_fragmented_partition_columns.year_partition_column
+                )
+
+                year_fragment_col = YearFragmentColumn(
+                    column_name=proto_athena_source.time_fragmented_partition_columns.year_partition_column.column_name,
+                    representation=year_rep,
+                )
+
+                month_fragment_col = None
+                if proto_athena_source.time_fragmented_partition_columns.HasField(
+                    "month_partition_column"
+                ):
+                    month_rep = AthenaSource._extract_column_representation(
+                        proto_athena_source.time_fragmented_partition_columns.month_partition_column
+                    )
+
+                    month_fragment_col = MonthFragmentColumn(
+                        column_name=proto_athena_source.time_fragmented_partition_columns.month_partition_column.column_name,
+                        representation=month_rep,
+                    )
+
+                day_fragment_col = None
+                if proto_athena_source.time_fragmented_partition_columns.HasField(
+                    "day_partition_column"
+                ):
+                    day_rep = AthenaSource._extract_column_representation(
+                        proto_athena_source.time_fragmented_partition_columns.day_partition_column
+                    )
+
+                    day_fragment_col = DayFragmentColumn(
+                        column_name=proto_athena_source.time_fragmented_partition_columns.day_partition_column.column_name,
+                        representation=day_rep,
+                    )
+
+                res = TimeFragmentedPartitionColumns(
+                    year_partition_column=year_fragment_col,
+                    month_partition_column=month_fragment_col,
+                    day_partition_column=day_fragment_col,
+                )
+            else:
+                raise ValueError(
+                    f"Got an unsupported time partition column: {proto_athena_source}"
+                )
+        else:
+            res = None
+
+        return res
+
+    @staticmethod
+    def _extract_column_representation(proto_fragment_column):
+        column_rep = proto_fragment_column.WhichOneof("column_representation")
+
+        if column_rep == "numeric_column_representation":
+            res = ColumnRepresentation.NumericColumnRepresentation
+        elif column_rep == "textual_column_representation":
+            res = ColumnRepresentation.TextualColumnRepresentation
+        else:
+            raise ValueError(f"Unsupported ColumnRepresentation: {column_rep}")
+
+        return res

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

@@ -3,15 +3,13 @@ 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 (
     BigquerySource as ProtoBigquerySource,
 )
 from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
     DataSourceSpec as ProtoDataSourceSpec,
 )
 from qwak.exceptions import QwakException
-from qwak.feature_store.data_sources.batch_sources._batch import BaseBatchSource
+from qwak.feature_store.data_sources.batch._batch import BaseBatchSource
 
 
 @dataclass
@@ -25,6 +23,7 @@ class BigQuerySource(BaseBatchSource):
     views_enabled: Optional[bool] = False
     materialization_dataset: Optional[str] = None
     materialization_project: Optional[str] = None
+    repository: Optional[str] = None
 
     def __post_init__(self):
         if not (self.sql or (self.table and self.dataset)):
@@ -59,8 +58,9 @@ class BigQuerySource(BaseBatchSource):
             materialization_dataset=bigquery.materialization_dataset,
         )
 
-    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,
@@ -76,5 +76,5 @@ class BigQuerySource(BaseBatchSource):
                     materialization_dataset=self.materialization_dataset,
                     materialization_project=self.materialization_project,
                 ),
-            )
+            ),
         )

qwak/feature_store/data_sources/batch/clickhouse.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
+    BatchSource as ProtoBatchSource,
+    ClickhouseSource as ProtoClickhouseSource,
+)
+from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
+    DataSourceSpec as ProtoDataSourceSpec,
+)
+from qwak.exceptions import QwakException
+from qwak.feature_store.data_sources.batch._batch import BaseBatchSource
+
+
+@dataclass
+class ClickhouseSource(BaseBatchSource):
+    username_secret_name: Optional[str] = None
+    password_secret_name: Optional[str] = None
+    url: Optional[str] = None
+    db_table: Optional[str] = None
+    query: Optional[str] = None
+    repository: Optional[str] = None
+
+    def __post_init__(self):
+        self._validate()
+
+    def _validate(self):
+        if not (bool(self.db_table) ^ bool(self.query)):
+            raise QwakException("Only one of query and db_table must be set")
+
+    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,
+                date_created_column=self.date_created_column,
+                clickhouseSource=ProtoClickhouseSource(
+                    url=self.url,
+                    username_secret_name=self.username_secret_name,
+                    password_secret_name=self.password_secret_name,
+                    table=self.db_table,
+                    sql=self.query,
+                ),
+            ),
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoBatchSource):
+        clickhouse: ProtoClickhouseSource = proto.clickhouseSource
+        return cls(
+            name=proto.name,
+            date_created_column=proto.date_created_column,
+            description=proto.description,
+            url=clickhouse.url,
+            username_secret_name=clickhouse.username_secret_name,
+            password_secret_name=clickhouse.password_secret_name,
+            db_table=clickhouse.table,
+            query=clickhouse.sql,
+        )

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

@@ -3,14 +3,16 @@ from typing import Optional
 
 from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
+    CsvSource as ProtoCsvSource,
 )
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import CsvSource as ProtoCsvSource
 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_sources.filesystem_config import (
+from qwak.feature_store.data_sources.batch._batch import BaseBatchSource
+from qwak.feature_store.data_sources.batch.filesystem.base_config import (
     FileSystemConfiguration,
+)
+from qwak.feature_store.data_sources.batch.filesystem.utils import (
     get_fs_config_from_proto,
 )
 
@@ -21,6 +23,7 @@ class CsvSource(BaseBatchSource):
     quote_character: str = '"'
     escape_character: str = '"'
     filesystem_configuration: Optional[FileSystemConfiguration] = None
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -38,12 +41,13 @@ class CsvSource(BaseBatchSource):
             filesystem_configuration=fs_conf,
         )
 
-    def _to_proto(self):
+    def _to_proto(self, artifact_url: Optional[str] = None):
         fs_conf = None
         if self.filesystem_configuration:
             fs_conf = self.filesystem_configuration._to_proto()
 
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -54,5 +58,5 @@ class CsvSource(BaseBatchSource):
                     escape_character=self.escape_character,
                     filesystem_configuration=fs_conf,
                 ),
-            )
+            ),
         )

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

@@ -3,14 +3,12 @@ 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 (
     ElasticsearchSource as ProtoElasticsearchSource,
 )
 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
@@ -23,6 +21,7 @@ class ElasticSearchSource(BaseBatchSource):
     password_secret_name: Optional[str] = None
     exclude_fields: Optional[str] = None
     parse_dates: Optional[bool] = True
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -41,8 +40,9 @@ class ElasticSearchSource(BaseBatchSource):
             parse_dates=elastic.parse_dates,
         )
 
-    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,
@@ -57,5 +57,5 @@ class ElasticSearchSource(BaseBatchSource):
                     exclude_fields=self.exclude_fields,
                     parse_dates=self.parse_dates,
                 ),
-            )
+            ),
         )