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

qwak/feature_store/data_sources/{batch_sources/filesystem_config.py → batch/filesystem/aws.py}

@@ -1,34 +1,22 @@
-from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Optional
 
 from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     AnonymousS3Configuration as ProtoAnonymousS3Configuration,
-)
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
+    AwsS3AssumeRole as ProtoAwsS3AssumeRole,
     AwsS3FileSystemConfiguration as ProtoAwsS3FileSystemConfiguration,
-)
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
-    FileSystemConfiguration as FileSystemConfigurationProto,
+    FileSystemConfiguration as ProtoFileSystemConfiguration,
 )
 from qwak.exceptions import QwakException
-
-
-@dataclass
-class FileSystemConfiguration(ABC):
-    @abstractmethod
-    def _to_proto(self):
-        pass
-
-    @abstractmethod
-    def _from_proto(self, proto):
-        pass
+from qwak.feature_store.data_sources.batch.filesystem.base_config import (
+    FileSystemConfiguration,
+)
 
 
 @dataclass
 class AnonymousS3Configuration(FileSystemConfiguration):
     def _to_proto(self):
-        return FileSystemConfigurationProto(
+        return ProtoFileSystemConfiguration(
             aws_s3_anonymous=ProtoAnonymousS3Configuration()
         )
 
@@ -37,6 +25,29 @@ class AnonymousS3Configuration(FileSystemConfiguration):
         return cls()
 
 
+@dataclass
+class AwsS3AssumeRoleFileSystemConfiguration(FileSystemConfiguration):
+    role_arn: str
+
+    def __post_init__(self):
+        self._validate()
+
+    def _validate(self):
+        if not self.role_arn:
+            raise QwakException("`role_arn` field is mandatory")
+
+    def _to_proto(self) -> ProtoAwsS3FileSystemConfiguration:
+        return ProtoFileSystemConfiguration(
+            aws_s3_assume_role_configuration=ProtoAwsS3AssumeRole(
+                role_arn=self.role_arn
+            )
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoAwsS3AssumeRole):
+        return cls(proto.role_arn)
+
+
 @dataclass
 class AwsS3FileSystemConfiguration(FileSystemConfiguration):
     access_key_secret_name: str
@@ -57,7 +68,7 @@ class AwsS3FileSystemConfiguration(FileSystemConfiguration):
             raise QwakException(error_msg.format(field="bucket"))
 
     def _to_proto(self):
-        return FileSystemConfigurationProto(
+        return ProtoFileSystemConfiguration(
             aws_s3_configuration=ProtoAwsS3FileSystemConfiguration(
                 access_key_secret_name=self.access_key_secret_name,
                 secret_key_secret_name=self.secret_key_secret_name,
@@ -74,22 +85,3 @@ class AwsS3FileSystemConfiguration(FileSystemConfiguration):
             bucket=proto.bucket,
             session_token_secret_name=proto.session_token_secret_name,
         )
-
-
-def get_fs_config_from_proto(filesystem_conf):
-    if not filesystem_conf:
-        return {}
-
-    fs_conf_type = filesystem_conf.WhichOneof("type")
-
-    if not fs_conf_type:
-        return {}
-
-    fs_conf = getattr(filesystem_conf, fs_conf_type)
-    if isinstance(fs_conf, ProtoAwsS3FileSystemConfiguration):
-        return AwsS3FileSystemConfiguration._from_proto(fs_conf)
-
-    elif isinstance(fs_conf, ProtoAnonymousS3Configuration):
-        return AnonymousS3Configuration._from_proto(fs_conf)
-    else:
-        raise QwakException(f"Unsupported FileSystemConfiguration: {fs_conf_type}")

qwak/feature_store/data_sources/batch/filesystem/base_config.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass
+class FileSystemConfiguration(ABC):
+    @abstractmethod
+    def _to_proto(self):
+        pass
+
+    @abstractmethod
+    def _from_proto(self, proto):
+        pass

qwak/feature_store/data_sources/batch/filesystem/gcp.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+
+from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
+    FileSystemConfiguration as ProtoFileSystemConfiguration,
+    GcsServiceAccountImpersonation as ProtoGcsServiceAccountImpersonation,
+    GcsUnauthenticated as ProtoGcsUnauthenticated,
+)
+from qwak.exceptions import QwakException
+from qwak.feature_store.data_sources.batch.filesystem.base_config import (
+    FileSystemConfiguration,
+)
+
+
+@dataclass
+class GcpGcsServiceAccountImpersonation(FileSystemConfiguration):
+    service_account_user: str
+
+    def __post_init__(self):
+        self._validate()
+
+    def _validate(self):
+        if not self.service_account_user or not self.service_account_user.strip():
+            raise QwakException(
+                "Service account user is mandatory for GCS service account impersonation, blanks are invalid"
+            )
+
+    def _to_proto(self):
+        return ProtoFileSystemConfiguration(
+            gcs_service_account_impersonation=ProtoGcsServiceAccountImpersonation(
+                service_account_user=self.service_account_user
+            )
+        )
+
+    @classmethod
+    def _from_proto(cls, proto):
+        return GcpGcsServiceAccountImpersonation(
+            service_account_user=proto.service_account_user
+        )
+
+
+@dataclass
+class GcpGcsUnauthenticated(FileSystemConfiguration):
+    def _to_proto(self):
+        return ProtoFileSystemConfiguration(
+            gcs_unauthenticated=ProtoGcsUnauthenticated()
+        )
+
+    @classmethod
+    def _from_proto(cls, proto):
+        return GcpGcsUnauthenticated()

qwak/feature_store/data_sources/batch/filesystem/utils.py

@@ -0,0 +1,51 @@
+from typing import Optional
+
+from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
+    AnonymousS3Configuration as ProtoAnonymousS3Configuration,
+    AwsS3AssumeRole as ProtoAwsS3AssumeRole,
+    AwsS3FileSystemConfiguration as ProtoAwsS3FileSystemConfiguration,
+    GcsServiceAccountImpersonation as ProtoGcsServiceAccountImpersonation,
+    GcsUnauthenticated as ProtoGcsUnauthenticated,
+)
+from qwak.exceptions import QwakException
+from qwak.feature_store.data_sources.batch.filesystem.aws import (
+    AnonymousS3Configuration,
+    AwsS3AssumeRoleFileSystemConfiguration,
+    AwsS3FileSystemConfiguration,
+)
+from qwak.feature_store.data_sources.batch.filesystem.base_config import (
+    FileSystemConfiguration,
+)
+from qwak.feature_store.data_sources.batch.filesystem.gcp import (
+    GcpGcsServiceAccountImpersonation,
+    GcpGcsUnauthenticated,
+)
+
+
+def get_fs_config_from_proto(filesystem_conf) -> Optional[FileSystemConfiguration]:
+    if not filesystem_conf:
+        return None
+
+    fs_conf_type = filesystem_conf.WhichOneof("type")
+
+    if not fs_conf_type:
+        return None
+
+    fs_conf = getattr(filesystem_conf, fs_conf_type)
+    if isinstance(fs_conf, ProtoAwsS3FileSystemConfiguration):
+        return AwsS3FileSystemConfiguration._from_proto(fs_conf)
+
+    elif isinstance(fs_conf, ProtoAnonymousS3Configuration):
+        return AnonymousS3Configuration._from_proto(fs_conf)
+
+    elif isinstance(fs_conf, ProtoAwsS3AssumeRole):
+        return AwsS3AssumeRoleFileSystemConfiguration._from_proto(fs_conf)
+
+    elif isinstance(fs_conf, ProtoGcsServiceAccountImpersonation):
+        return GcpGcsServiceAccountImpersonation._from_proto(fs_conf)
+
+    elif isinstance(fs_conf, ProtoGcsUnauthenticated):
+        return GcpGcsUnauthenticated._from_proto(fs_conf)
+
+    else:
+        raise QwakException(f"Unsupported FileSystemConfiguration: {fs_conf_type}")

qwak/feature_store/data_sources/{batch_sources → batch}/mongodb.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 (
     MongoSource as ProtoMongoSource,
 )
 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 MongoDbSource(BaseBatchSource):
     collection: str
     connection_params: str
     protocol: str = "mongodb"
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -38,8 +38,9 @@ class MongoDbSource(BaseBatchSource):
             protocol=mongo.protocol,
         )
 
-    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 MongoDbSource(BaseBatchSource):
                     connection_params=self.connection_params,
                     protocol=self.protocol,
                 ),
-            )
+            ),
         )

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

@@ -1,18 +1,15 @@
 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 (
     JdbcSource as ProtoJdbcSource,
-)
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     MysqlSource as ProtoMysqlSource,
 )
 from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
     DataSourceSpec as ProtoDataSourceSpec,
 )
-from qwak.feature_store.data_sources.batch_sources._jdbc import JdbcSource
+from qwak.feature_store.data_sources.batch._jdbc import JdbcSource
 
 
 @dataclass
@@ -31,8 +28,9 @@ class MysqlSource(JdbcSource):
             query=mysql.query,
         )
 
-    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,
@@ -45,5 +43,5 @@ class MysqlSource(JdbcSource):
                     query=self.query,
                     mysqlSource=ProtoMysqlSource(),
                 ),
-            )
+            ),
         )

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

@@ -1,18 +1,19 @@
 from dataclasses import dataclass, field
+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 (
     ParquetSource as ProtoParquetSource,
 )
 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 (
-    AnonymousS3Configuration,
+from qwak.feature_store.data_sources import AnonymousS3Configuration
+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,
 )
 
@@ -23,6 +24,7 @@ class ParquetSource(BaseBatchSource):
     filesystem_configuration: FileSystemConfiguration = field(
         default_factory=lambda: AnonymousS3Configuration()
     )
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -37,12 +39,13 @@ class ParquetSource(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,
@@ -50,5 +53,5 @@ class ParquetSource(BaseBatchSource):
                 parquetSource=ProtoParquetSource(
                     path=self.path, filesystem_configuration=fs_conf
                 ),
-            )
+            ),
         )

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

@@ -1,24 +1,22 @@
 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 (
     JdbcSource as ProtoJdbcSource,
-)
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     PostgresqlSource as ProtoPostgresqlSource,
 )
 from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
     DataSourceSpec as ProtoDataSourceSpec,
 )
-from qwak.feature_store.data_sources.batch_sources._jdbc import JdbcSource
+from qwak.feature_store.data_sources.batch._jdbc import JdbcSource
 
 
 @dataclass
 class PostgresSource(JdbcSource):
-    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,
@@ -31,7 +29,7 @@ class PostgresSource(JdbcSource):
                     query=self.query,
                     postgresqlSource=ProtoPostgresqlSource(),
                 ),
-            )
+            ),
         )
 
     @classmethod

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

@@ -3,18 +3,14 @@ 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 (
     JdbcSource as ProtoJdbcSource,
-)
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     RedshiftSource as ProtoRedshiftSource,
 )
 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._jdbc import JdbcSource
+from qwak.feature_store.data_sources.batch._jdbc import JdbcSource
 
 
 @dataclass
@@ -27,6 +23,8 @@ class RedshiftSource(JdbcSource):
     secret_access_key: Optional[str] = None
     query_group: Optional[str] = "_qwak_featurestore"
 
+    repository: Optional[str] = None
+
     def _validate(self):
         authentication_methods = sum(
             [
@@ -46,8 +44,9 @@ class RedshiftSource(JdbcSource):
                 "or user id secret name and password secret name"
             )
 
-    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,
@@ -65,7 +64,7 @@ class RedshiftSource(JdbcSource):
                         secret_access_key=self.secret_access_key,
                     ),
                 ),
-            )
+            ),
         )
 
     @classmethod

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

@@ -1,47 +1,71 @@
 from dataclasses import dataclass
 from typing import Optional
+import warnings
 
 from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
-)
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     SnowflakeSource as ProtoSnowflakeSource,
 )
 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
+
+warnings.simplefilter("once", DeprecationWarning)
 
 
 @dataclass
 class SnowflakeSource(BaseBatchSource):
     host: str
     username_secret_name: str
-    password_secret_name: str
     database: str
     schema: str
     warehouse: str
+    password_secret_name: Optional[str] = None
+    pem_private_key_secret_name: Optional[str] = None
     table: Optional[str] = None
     query: Optional[str] = None
+    repository: Optional[str] = None
 
     def __post_init__(self):
         self._validate()
 
     def _validate(self):
-        if self.table and self.query:
-            raise QwakException("Only one of query and table may be set")
-        if not self.table and not self.query:
-            raise QwakException("One of table or query must be set")
-
         if not self.username_secret_name:
             raise QwakException("username_secret_name must be set!")
 
-        if not self.password_secret_name:
-            raise QwakException("password_secret_name must be set!")
+        if not self.database:
+            raise QwakException("database must be set!")
+
+        if not self.schema:
+            raise QwakException("schema must be set!")
 
-    def _to_proto(self):
+        if self.password_secret_name:
+            warnings.warn(
+                "Snowflake basic authentication is deprecated and should not be used. Use key-pair authentication instead.",
+                DeprecationWarning,
+            )
+
+        no_unique_source_exception_message = "Only one of query or table may be set"
+        has_table = bool(self.table)
+        has_query = bool(self.query)
+        no_source_set = not (has_table or has_query)
+        both_source_set = has_table and has_query
+        if no_source_set or both_source_set:
+            raise QwakException(no_unique_source_exception_message)
+
+        no_unique_auth_exception_message = "Exactly one of 'password_secret_name' or 'pem_private_key_secret_name' must be set"
+        has_password_secret = bool(self.password_secret_name)
+        has_pem_private_key_secret = bool(self.pem_private_key_secret_name)
+        no_auth_set = not (has_password_secret or has_pem_private_key_secret)
+        both_auth_set = has_password_secret and has_pem_private_key_secret
+        if no_auth_set or both_auth_set:
+            raise QwakException(no_unique_auth_exception_message)
+
+    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,
@@ -50,13 +74,14 @@ class SnowflakeSource(BaseBatchSource):
                     host=self.host,
                     username_secret_name=self.username_secret_name,
                     password_secret_name=self.password_secret_name,
+                    pem_private_key_secret_name=self.pem_private_key_secret_name,
                     database=self.database,
                     schema=self.schema,
                     warehouse=self.warehouse,
                     table=self.table,
                     query=self.query,
                 ),
-            )
+            ),
         )
 
     @classmethod
@@ -69,6 +94,7 @@ class SnowflakeSource(BaseBatchSource):
             host=snowflake.host,
             username_secret_name=snowflake.username_secret_name,
             password_secret_name=snowflake.password_secret_name,
+            pem_private_key_secret_name=snowflake.pem_private_key_secret_name,
             database=snowflake.database,
             schema=snowflake.schema,
             warehouse=snowflake.warehouse,

qwak/feature_store/data_sources/batch/unity_catalog.py

@@ -0,0 +1,107 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
+    BatchSource as ProtoBatchSource,
+    CatalogSource as ProtoCatalogSource,
+    UnityCatalogSource as ProtoUnityCatalogSource,
+)
+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
+
+RESERVED_CATALOG_NAME = "qwak_catalog"
+
+
+@dataclass
+class UnityCatalogSource(BaseBatchSource):
+    """Unity Catalog batch data source.
+
+    Attributes:
+        uri (str): Connection uri to Unity Catalog i.e https://<workspace-instance-name>.<domain-name>/api/2.1/unity-catalog
+
+        schema (str): The Unity Catalog schema name where the data is located.
+
+        catalog (str): The Unity Catalog catalog name.
+
+        personal_access_token_secret_name (Optional[str]): The name of the secret
+            containing the personal access token for authentication. e.g "my-pat-secret"
+
+        table (Optional[str]): The table name within the specified catalog and schema.
+            Either 'table' or 'query' must be provided, but not both. e.g "my_table"
+
+        query (Optional[str]): A SQL query to execute against the Unity Catalog.
+            Either 'query' or 'table' must be provided, but not both. e.g "select * from {catalog}.{schema}.{table}"
+
+        repository (Optional[str]): The data source group.
+
+    """
+
+    uri: str
+    schema: str
+    catalog: str
+    personal_access_token_secret_name: Optional[str] = None
+    table: Optional[str] = None
+    query: Optional[str] = None
+    repository: Optional[str] = None
+
+    def __post_init__(self):
+        self._validate()
+
+    def _validate(self):
+        if not self.uri:
+            raise QwakException("`uri` cannot be empty.")
+
+        if not self.schema:
+            raise QwakException("`schema` cannot be empty.")
+
+        if not self.catalog:
+            raise QwakException("`catalog` cannot be empty.")
+
+        if not self.personal_access_token_secret_name:
+            raise QwakException("`personal_access_token_secret_name` cannot be empty.")
+
+        if self.catalog == RESERVED_CATALOG_NAME:
+            raise QwakException(
+                f"`catalog` can not be set to `{RESERVED_CATALOG_NAME}`.",
+            )
+
+        if bool(self.query) == bool(self.table):
+            raise QwakException("Either `query` or `table` must be provided.")
+
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
+        return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
+            batch_source=ProtoBatchSource(
+                name=self.name,
+                description=self.description,
+                date_created_column=self.date_created_column,
+                catalogSource=ProtoCatalogSource(
+                    url=self.uri,
+                    query=self.query,
+                    table=self.table,
+                    unityCatalogSource=ProtoUnityCatalogSource(
+                        catalog=self.catalog,
+                        schema=self.schema,
+                        personal_access_token_secret_name=self.personal_access_token_secret_name,
+                    ),
+                ),
+            ),
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoBatchSource) -> "UnityCatalogSource":
+        catalog_source = proto.catalogSource
+        return cls(
+            name=proto.name,
+            date_created_column=proto.date_created_column,
+            description=proto.description,
+            uri=catalog_source.url,
+            schema=catalog_source.unityCatalogSource.schema,
+            catalog=catalog_source.unityCatalogSource.catalog,
+            personal_access_token_secret_name=catalog_source.unityCatalogSource.personal_access_token_secret_name,
+            table=catalog_source.table,
+            query=catalog_source.query,
+        )