frogml 1.1.115__py3-none-any.whl → 1.2.19__py3-none-any.whl

frogml/feature_store/data_sources/batch/_jdbc.py

@@ -1,23 +1,22 @@
 from abc import ABC
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
+from pydantic import model_validator
 
 from frogml.core.exceptions import FrogmlException
 from frogml.feature_store.data_sources.batch._batch import BaseBatchSource
 
 
-@dataclass
 class JdbcSource(BaseBatchSource, ABC):
     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):
+    @model_validator(mode="after")
+    def __validate_jdbc(self) -> Self:
         if not (bool(self.db_table) ^ bool(self.query)):
             raise FrogmlException("Only one of query and db_table must be set")
+
+        return self

frogml/feature_store/data_sources/batch/athena.py

@@ -1,5 +1,6 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
+from pydantic import model_validator
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     AthenaSourceV1 as ProtoAthenaSourceV1,
@@ -45,18 +46,15 @@ from frogml.core.feature_store.data_sources.time_partition_columns import (
 )
 
 
-@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()
 
+    @model_validator(mode="after")
+    def __validate_athena(self) -> Self:
         if self.aws_authentication is None:
             raise AttributeError("`aws_authentication` must be set")
 
@@ -67,6 +65,7 @@ class AthenaSource(JdbcSource):
             raise AttributeError(
                 "At least one of: `s3_output_location`, `workgroup` must be set"
             )
+        return self
 
     def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         proto_date_partition_columns = None
@@ -122,7 +121,7 @@ class AthenaSource(JdbcSource):
         )
 
     @classmethod
-    def _from_proto(cls, proto: ProtoBatchSource):
+    def _from_proto(cls, proto: ProtoBatchSource) -> Self:
         proto_jdbc_source: ProtoJdbcSource = proto.jdbcSource
         proto_athena_source: ProtoAthenaSourceV1 = proto_jdbc_source.athenaSource
 

frogml/feature_store/data_sources/batch/big_query.py

@@ -1,5 +1,7 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
+
+from pydantic import model_validator
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -14,7 +16,6 @@ from frogml.core.exceptions import FrogmlException
 from frogml.feature_store.data_sources.batch._batch import BaseBatchSource
 
 
-@dataclass
 class BigQuerySource(BaseBatchSource):
     credentials_secret_name: str
     project: Optional[str] = None
@@ -25,9 +26,9 @@ 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):
+    @model_validator(mode="after")
+    def __validate_bigquery(self) -> Self:
         if not (self.sql or (self.table and self.dataset)):
             raise FrogmlException("Either SQL or (table and dataset) must be provided.")
 
@@ -42,8 +43,10 @@ class BigQuerySource(BaseBatchSource):
                 )
             self.views_enabled = True
 
+        return self
+
     @classmethod
-    def _from_proto(cls, proto):
+    def _from_proto(cls, proto) -> Self:
         bigquery = proto.bigquerySource
         return cls(
             name=proto.name,

frogml/feature_store/data_sources/batch/clickhouse.py

@@ -1,5 +1,7 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
+
+from pydantic import model_validator
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -14,22 +16,20 @@ from frogml.core.exceptions import FrogmlException
 from frogml.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):
+    @model_validator(mode="after")
+    def __validate_clickhouse(self) -> Self:
         if not (bool(self.db_table) ^ bool(self.query)):
             raise FrogmlException("Only one of query and db_table must be set")
 
+        return self
+
     def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
             data_source_repository_name=self.repository,
@@ -48,7 +48,7 @@ class ClickhouseSource(BaseBatchSource):
         )
 
     @classmethod
-    def _from_proto(cls, proto: ProtoBatchSource):
+    def _from_proto(cls, proto: ProtoBatchSource) -> Self:
         clickhouse: ProtoClickhouseSource = proto.clickhouseSource
         return cls(
             name=proto.name,

frogml/feature_store/data_sources/batch/csv.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -19,16 +19,14 @@ from frogml.feature_store.data_sources.batch.filesystem.utils import (
 )
 
 
-@dataclass
 class CsvSource(BaseBatchSource):
     path: str
     quote_character: str = '"'
     escape_character: str = '"'
     filesystem_configuration: Optional[FileSystemConfiguration] = None
-    repository: Optional[str] = None
 
     @classmethod
-    def _from_proto(cls, proto):
+    def _from_proto(cls, proto) -> Self:
         csv = proto.csvSource
 
         fs_conf = get_fs_config_from_proto(csv.filesystem_configuration)

frogml/feature_store/data_sources/batch/elastic_search.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -13,7 +13,6 @@ from frogml._proto.qwak.feature_store.sources.data_source_pb2 import (
 from frogml.feature_store.data_sources.batch._batch import BaseBatchSource
 
 
-@dataclass
 class ElasticSearchSource(BaseBatchSource):
     nodes: str
     resource: str
@@ -23,10 +22,9 @@ 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):
+    def _from_proto(cls, proto) -> Self:
         elastic = proto.elasticsearchSource
         return cls(
             name=proto.name,

frogml/feature_store/data_sources/batch/filesystem/aws.py

@@ -1,5 +1,7 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
+
+from pydantic import model_validator
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     AnonymousS3Configuration as ProtoAnonymousS3Configuration,
@@ -19,30 +21,28 @@ from frogml.feature_store.data_sources.batch.filesystem.base_config import (
 )
 
 
-@dataclass
 class AnonymousS3Configuration(FileSystemConfiguration):
-    def _to_proto(self):
+    def _to_proto(self) -> ProtoFileSystemConfiguration:
         return ProtoFileSystemConfiguration(
             aws_s3_anonymous=ProtoAnonymousS3Configuration()
         )
 
     @classmethod
-    def _from_proto(cls, proto):
+    def _from_proto(cls, proto: ProtoAnonymousS3Configuration) -> Self:
         return cls()
 
 
-@dataclass
 class AwsS3AssumeRoleFileSystemConfiguration(FileSystemConfiguration):
     role_arn: str
 
-    def __post_init__(self):
-        self._validate()
-
-    def _validate(self):
+    @model_validator(mode="after")
+    def __validate_role_arn(self) -> Self:
         if not self.role_arn:
             raise FrogmlException("`role_arn` field is mandatory")
 
-    def _to_proto(self) -> ProtoAwsS3FileSystemConfiguration:
+        return self
+
+    def _to_proto(self) -> ProtoFileSystemConfiguration:
         return ProtoFileSystemConfiguration(
             aws_s3_assume_role_configuration=ProtoAwsS3AssumeRole(
                 role_arn=self.role_arn
@@ -50,21 +50,18 @@ class AwsS3AssumeRoleFileSystemConfiguration(FileSystemConfiguration):
         )
 
     @classmethod
-    def _from_proto(cls, proto: ProtoAwsS3AssumeRole):
-        return cls(proto.role_arn)
+    def _from_proto(cls, proto: ProtoAwsS3AssumeRole) -> Self:
+        return cls(role_arn=proto.role_arn)
 
 
-@dataclass
 class AwsS3FileSystemConfiguration(FileSystemConfiguration):
-    access_key_secret_name: str
-    secret_key_secret_name: str
-    bucket: str
+    access_key_secret_name: Optional[str] = None
+    secret_key_secret_name: Optional[str] = None
+    bucket: Optional[str] = None
     session_token_secret_name: Optional[str] = ""
 
-    def __post_init__(self):
-        self._validate()
-
-    def _validate(self):
+    @model_validator(mode="after")
+    def __validate_fs_configuration(self) -> Self:
         error_msg = "{field} field is mandatory"
         if not self.access_key_secret_name:
             raise FrogmlException(error_msg.format(field="access_key"))
@@ -73,7 +70,9 @@ class AwsS3FileSystemConfiguration(FileSystemConfiguration):
         if not self.bucket:
             raise FrogmlException(error_msg.format(field="bucket"))
 
-    def _to_proto(self):
+        return self
+
+    def _to_proto(self) -> ProtoFileSystemConfiguration:
         return ProtoFileSystemConfiguration(
             aws_s3_configuration=ProtoAwsS3FileSystemConfiguration(
                 access_key_secret_name=self.access_key_secret_name,
@@ -84,7 +83,7 @@ class AwsS3FileSystemConfiguration(FileSystemConfiguration):
         )
 
     @classmethod
-    def _from_proto(cls, proto):
+    def _from_proto(cls, proto: ProtoAwsS3FileSystemConfiguration) -> Self:
         return AwsS3FileSystemConfiguration(
             access_key_secret_name=proto.access_key_secret_name,
             secret_key_secret_name=proto.secret_key_secret_name,

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

@@ -1,9 +1,9 @@
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
 
+from pydantic import BaseModel
 
-@dataclass
-class FileSystemConfiguration(ABC):
+
+class FileSystemConfiguration(BaseModel, ABC):
     @abstractmethod
     def _to_proto(self):
         pass

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

@@ -1,5 +1,4 @@
-from dataclasses import dataclass
-
+from pydantic import model_validator
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     FileSystemConfiguration as ProtoFileSystemConfiguration,
 )
@@ -15,18 +14,16 @@ from frogml.feature_store.data_sources.batch.filesystem.base_config import (
 )
 
 
-@dataclass
 class GcpGcsServiceAccountImpersonation(FileSystemConfiguration):
     service_account_user: str
 
-    def __post_init__(self):
-        self._validate()
-
+    @model_validator(mode="after")
     def _validate(self):
         if not self.service_account_user or not self.service_account_user.strip():
             raise FrogmlException(
                 "Service account user is mandatory for GCS service account impersonation, blanks are invalid"
             )
+        return self
 
     def _to_proto(self):
         return ProtoFileSystemConfiguration(
@@ -42,7 +39,6 @@ class GcpGcsServiceAccountImpersonation(FileSystemConfiguration):
         )
 
 
-@dataclass
 class GcpGcsUnauthenticated(FileSystemConfiguration):
     def _to_proto(self):
         return ProtoFileSystemConfiguration(

frogml/feature_store/data_sources/batch/mongodb.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -13,7 +13,6 @@ from frogml._proto.qwak.feature_store.sources.data_source_pb2 import (
 from frogml.feature_store.data_sources.batch._batch import BaseBatchSource
 
 
-@dataclass
 class MongoDbSource(BaseBatchSource):
     hosts: str
     username_secret_name: str
@@ -22,10 +21,9 @@ class MongoDbSource(BaseBatchSource):
     collection: str
     connection_params: str
     protocol: str = "mongodb"
-    repository: Optional[str] = None
 
     @classmethod
-    def _from_proto(cls, proto):
+    def _from_proto(cls, proto) -> Self:
         mongo = proto.mongoSource
         return cls(
             name=proto.name,

frogml/feature_store/data_sources/batch/mysql.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -16,10 +16,9 @@ from frogml._proto.qwak.feature_store.sources.data_source_pb2 import (
 from frogml.feature_store.data_sources.batch._jdbc import JdbcSource
 
 
-@dataclass
 class MysqlSource(JdbcSource):
     @classmethod
-    def _from_proto(cls, proto):
+    def _from_proto(cls, proto) -> Self:
         mysql = proto.jdbcSource
         return cls(
             name=proto.name,

frogml/feature_store/data_sources/batch/parquet.py

@@ -1,5 +1,7 @@
-from dataclasses import dataclass, field
 from typing import Optional
+from typing_extensions import Self
+
+from pydantic import Field
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -20,16 +22,14 @@ from frogml.feature_store.data_sources.batch.filesystem.utils import (
 )
 
 
-@dataclass
 class ParquetSource(BaseBatchSource):
     path: str
-    filesystem_configuration: FileSystemConfiguration = field(
-        default_factory=lambda: AnonymousS3Configuration()
+    filesystem_configuration: Optional[FileSystemConfiguration] = Field(
+        default_factory=AnonymousS3Configuration
     )
-    repository: Optional[str] = None
 
     @classmethod
-    def _from_proto(cls, proto):
+    def _from_proto(cls, proto: ProtoBatchSource) -> Self:
         parquet: ProtoParquetSource = proto.parquetSource
         fs_conf = get_fs_config_from_proto(parquet.filesystem_configuration)
 

frogml/feature_store/data_sources/batch/postgres.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -16,7 +16,6 @@ from frogml._proto.qwak.feature_store.sources.data_source_pb2 import (
 from frogml.feature_store.data_sources.batch._jdbc import JdbcSource
 
 
-@dataclass
 class PostgresSource(JdbcSource):
     def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
@@ -37,7 +36,7 @@ class PostgresSource(JdbcSource):
         )
 
     @classmethod
-    def _from_proto(cls, proto):
+    def _from_proto(cls, proto) -> Self:
         postgresql = proto.jdbcSource
         return cls(
             name=proto.name,

frogml/feature_store/data_sources/batch/redshift.py

@@ -1,5 +1,7 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
+
+from pydantic import model_validator
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -17,7 +19,6 @@ from frogml.core.exceptions import FrogmlException
 from frogml.feature_store.data_sources.batch._jdbc import JdbcSource
 
 
-@dataclass
 class RedshiftSource(JdbcSource):
     iam_role_arn: Optional[str] = None
     db_user: Optional[str] = None
@@ -27,9 +28,8 @@ class RedshiftSource(JdbcSource):
     secret_access_key: Optional[str] = None
     query_group: Optional[str] = "_qwak_featurestore"
 
-    repository: Optional[str] = None
-
-    def _validate(self):
+    @model_validator(mode="after")
+    def __validate_redshift(self) -> Self:
         authentication_methods = sum(
             [
                 bool(self.access_key and self.secret_access_key),
@@ -48,6 +48,8 @@ class RedshiftSource(JdbcSource):
                 "or user id secret name and password secret name"
             )
 
+        return self
+
     def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
             data_source_repository_name=self.repository,
@@ -72,7 +74,7 @@ class RedshiftSource(JdbcSource):
         )
 
     @classmethod
-    def _from_proto(cls, proto: ProtoBatchSource):
+    def _from_proto(cls, proto: ProtoBatchSource) -> Self:
         redshift = proto.jdbcSource
         return cls(
             name=proto.name,

frogml/feature_store/data_sources/batch/snowflake.py

@@ -1,6 +1,8 @@
 import warnings
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
+
+from pydantic import Field, model_validator
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -17,30 +19,32 @@ from frogml.feature_store.data_sources.batch._batch import BaseBatchSource
 warnings.simplefilter("once", DeprecationWarning)
 
 
-@dataclass
 class SnowflakeSource(BaseBatchSource):
     host: str
     username_secret_name: str
-    database: str
-    schema: str
+    database: Optional[str] = None
+    schema_name: Optional[str] = Field(
+        default=None, alias="schema", serialization_alias="schema"
+    )
     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()
+    @property
+    def schema(self) -> Optional[str]:
+        return self.schema_name
 
-    def _validate(self):
+    @model_validator(mode="after")
+    def __validate_snowflake(self) -> Self:
         if not self.username_secret_name:
             raise FrogmlException("username_secret_name must be set!")
 
         if not self.database:
             raise FrogmlException("database must be set!")
 
-        if not self.schema:
+        if not self.schema_name:
             raise FrogmlException("schema must be set!")
 
         if self.password_secret_name:
@@ -65,6 +69,8 @@ class SnowflakeSource(BaseBatchSource):
         if no_auth_set or both_auth_set:
             raise FrogmlException(no_unique_auth_exception_message)
 
+        return self
+
     def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
             data_source_repository_name=self.repository,
@@ -78,7 +84,7 @@ class SnowflakeSource(BaseBatchSource):
                     password_secret_name=self.password_secret_name,
                     pem_private_key_secret_name=self.pem_private_key_secret_name,
                     database=self.database,
-                    schema=self.schema,
+                    schema=self.schema_name,
                     warehouse=self.warehouse,
                     table=self.table,
                     query=self.query,
@@ -87,7 +93,7 @@ class SnowflakeSource(BaseBatchSource):
         )
 
     @classmethod
-    def _from_proto(cls, proto):
+    def _from_proto(cls, proto: ProtoBatchSource) -> Self:
         snowflake = proto.snowflakeSource
         return cls(
             name=proto.name,

frogml/feature_store/data_sources/batch/unity_catalog.py

@@ -1,5 +1,7 @@
-from dataclasses import dataclass
 from typing import Optional
+from typing_extensions import Self
+
+from pydantic import Field, model_validator
 
 from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -15,7 +17,6 @@ from frogml.feature_store.data_sources.batch._batch import BaseBatchSource
 RESERVED_CATALOG_NAME = "qwak_catalog"
 
 
-@dataclass
 class UnityCatalogSource(BaseBatchSource):
     """Unity Catalog batch data source.
 
@@ -40,21 +41,22 @@ class UnityCatalogSource(BaseBatchSource):
     """
 
     uri: str
-    schema: str
+    schema_name: str = Field(alias="schema", serialization_alias="schema")
     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()
+    @property
+    def schema(self) -> str:
+        return self.schema_name
 
-    def _validate(self):
+    @model_validator(mode="after")
+    def __validate_unity_catalog(self) -> Self:
         if not self.uri:
             raise FrogmlException("`uri` cannot be empty.")
 
-        if not self.schema:
+        if not self.schema_name:
             raise FrogmlException("`schema` cannot be empty.")
 
         if not self.catalog:
@@ -73,6 +75,8 @@ class UnityCatalogSource(BaseBatchSource):
         if bool(self.query) == bool(self.table):
             raise FrogmlException("Either `query` or `table` must be provided.")
 
+        return self
+
     def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
             data_source_repository_name=self.repository,
@@ -86,7 +90,7 @@ class UnityCatalogSource(BaseBatchSource):
                     table=self.table,
                     unityCatalogSource=ProtoUnityCatalogSource(
                         catalog=self.catalog,
-                        schema=self.schema,
+                        schema=self.schema_name,
                         personal_access_token_secret_name=self.personal_access_token_secret_name,
                     ),
                 ),
@@ -94,7 +98,7 @@ class UnityCatalogSource(BaseBatchSource):
         )
 
     @classmethod
-    def _from_proto(cls, proto: ProtoBatchSource) -> "UnityCatalogSource":
+    def _from_proto(cls, proto: ProtoBatchSource) -> Self:
         catalog_source = proto.catalogSource
         return cls(
             name=proto.name,