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

qwak/clients/project/client.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import grpc
 from _qwak_proto.qwak.projects.projects_pb2 import (
     CreateProjectRequest,
@@ -8,7 +6,6 @@ from _qwak_proto.qwak.projects.projects_pb2 import (
     ListProjectsRequest,
 )
 from _qwak_proto.qwak.projects.projects_pb2_grpc import ProjectsManagementServiceStub
-from _qwak_proto.qwak.projects.jfrog_project_spec_pb2 import ModelRepositoryJFrogSpec
 from dependency_injector.wiring import Provide, inject
 from qwak.exceptions import QwakException
 from qwak.inner.di_configuration import QwakContainer
@@ -34,16 +31,12 @@ class ProjectsManagementClient:
         self,
         project_name,
         project_description,
-        jfrog_project_key: Optional[str] = None,
     ):
         try:
             return self._projects_management_service.CreateProject(
                 CreateProjectRequest(
                     project_name=project_name,
                     project_description=project_description,
-                    jfrog_spec=ModelRepositoryJFrogSpec(
-                        jfrog_project_key=jfrog_project_key
-                    ),
                 )
             )
 

qwak/exceptions/__init__.py

@@ -1,6 +1,7 @@
 from .quiet_error import QuietError
 from .qwak_exception import QwakException
 from .qwak_general_build_exception import QwakGeneralBuildException
+from .qwak_grpc_address_exception import QwakGrpcAddressException
 from .qwak_http_exception import QwakHTTPException
 from .qwak_inference_exception import QwakInferenceException
 from .qwak_load_configuration_exception import LoadConfigurationException

qwak/exceptions/qwak_grpc_address_exception.py

@@ -0,0 +1,9 @@
+from typing import Union
+from urllib.parse import ParseResult
+
+from qwak.exceptions import QwakException
+
+
+class QwakGrpcAddressException(QwakException):
+    def __init__(self, details: str, grpc_address: Union[str, ParseResult]):
+        self.message = f"Not a valid gRPC address: '{grpc_address}'. Details: {details}"

qwak/feature_store/_common/artifact_utils.py

@@ -23,7 +23,7 @@ class ArtifactSpec:
     artifact_name: str
     root_module_path: Path
     artifact_object: Union["BaseTransformation", "Deserializer"]
-    callables: List[Callable] = field(default_factory=lambda: list())
+    callables: List[Callable] = field(default_factory=list)
     suffix: str = ""
 
 
@@ -48,10 +48,12 @@ class ArtifactsUploader:
         transformation_functions: Optional[List[Callable[..., Any]]] = (
             transformation.get_functions()
         )
-        if transformation_functions is not None and transformation_functions:
-            return ArtifactSpec(
-                artifact_name=featureset_name,
-                root_module_path=Path(__instance_module_path__).parent,
-                artifact_object=transformation,
-                callables=transformation_functions,
-            )
+        if not transformation_functions:
+            return None
+
+        return ArtifactSpec(
+            artifact_name=featureset_name,
+            root_module_path=Path(__instance_module_path__).parent,
+            artifact_object=transformation,
+            callables=transformation_functions,
+        )

qwak/feature_store/_common/packaging.py

@@ -6,16 +6,17 @@ import zipfile
 from io import BytesIO
 from pathlib import Path
 from shutil import make_archive
-from typing import TYPE_CHECKING, Callable, List, Union
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
 from uuid import uuid4
 
 import requests
-
 from qwak.clients.feature_store.management_client import FeatureRegistryClient
 from qwak.exceptions import QwakException
 from qwak.feature_store._common.value import (
     UPDATE_QWAK_SDK_WITH_FEATURE_STORE_EXTRA_MSG,
 )
+from requests import HTTPError, Response
+from urllib.parse import urlparse, ParseResult
 
 ZIP_FUNCTION_CONTENT_TYPE = "text/plain"
 
@@ -195,29 +196,47 @@ def upload_artifact(
     return uploaded_artifact_url
 
 
-def put_presigned_url(presign_url: str, content: bytes, content_type: str) -> None:
-    http_response = requests.put(
-        presign_url,
-        headers={"content-type": content_type},
-        data=content,
-        timeout=600,
-    )
-    if 400 <= http_response.status_code < 600:
-        raise QwakException(
-            f"Error while uploading user function to s3. status code: {http_response.status_code}"
+def put_files_content(
+    url: str,
+    content: bytes,
+    content_type: str,
+    extra_headers: Optional[Dict[str, str]] = None,
+) -> None:
+    """
+    uploads files content to s3 using presigned url or to jfrog artifactory using url and authentication headers
+    """
+    extra_headers = extra_headers or {}
+    try:
+        http_response: Response = requests.put(
+            url,
+            headers={"content-type": content_type, **extra_headers},
+            data=content,
+            timeout=300,
         )
+        http_response.raise_for_status()
+    except HTTPError as e:
+        raise QwakException(
+            f"Error while uploading artifact. status code: {e.response.status_code} with reason: {e.response.reason}"
+        ) from e
 
 
 def upload_to_s3(presign_url: str, in_memory_zip: bytes, content_type: str) -> str:
-    put_presigned_url(
-        presign_url=presign_url, content=in_memory_zip, content_type=content_type
+    parsed_url: ParseResult = urlparse(presign_url)
+    extra_headers: Dict[str, str] = (
+        {"x-ms-blob-type": "BlockBlob"}
+        if parsed_url.hostname.endswith(".blob.core.windows.net")
+        else None
+    )
+    put_files_content(
+        url=presign_url,
+        content=in_memory_zip,
+        content_type=content_type,
+        extra_headers=extra_headers,
     )
-    from urllib.parse import urlparse
 
-    parsed_url = urlparse(presign_url)
-    postfix = parsed_url.path
-    bucket = parsed_url.hostname.split(".s3.")[0]
-    s3_path = f"s3://{bucket}{postfix}"  # noqa: E231
+    postfix: str = parsed_url.path
+    bucket: str = parsed_url.hostname.split(".s3.")[0]
+    s3_path: str = f"s3://{bucket}{postfix}"  # noqa: E231
     return s3_path
 
 
@@ -227,8 +246,7 @@ def base_upload_function_code(
     feature_module_dir: Path,
 ) -> str:
     """
-    get presign url and prepare a zip file with the functions code
-    and it's module and upload it to clients' s3 bucket
+    get presign url and prepare a zip file with the functions code and it's module than uploads it to clients' s3 bucket
     Args:
         presign_url: presign url
         functions: udf functions

qwak/feature_store/_common/source_code_spec.py

@@ -1,48 +1,98 @@
 from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
 from typing import Optional
 
+from _qwak_proto.qwak.feature_store.v1.common.jfrog_artifact.jfrog_artifact_pb2 import (
+    JfrogArtifact as ProtoJfrogArtifact,
+)
 from _qwak_proto.qwak.feature_store.v1.common.source_code.source_code_pb2 import (
-    SourceCodeSpec as ProtoSourceCodeSpec,
     SourceCodeArtifact as ProtoSourceCodeArtifact,
+    SourceCodeSpec as ProtoSourceCodeSpec,
     ZipArtifact as ProtoZipArtifact,
 )
+from pydantic import BaseModel, Field, model_validator
 from qwak.exceptions import QwakException
+from typing_extensions import Self
+
+
+class JfrogArtifact(BaseModel):
+    path: str
+
+    @classmethod
+    def from_proto(cls, proto: ProtoJfrogArtifact) -> "JfrogArtifact":
+        return cls(path=proto.path)
+
+    def to_proto(self) -> ProtoJfrogArtifact:
+        return ProtoJfrogArtifact(path=self.path)
 
 
-class SourceCodeArtifact(ABC):
+class SourceCodeArtifact(BaseModel, ABC):
     @abstractmethod
     def _to_proto(self) -> ProtoSourceCodeArtifact:
         pass
 
 
-@dataclass
 class ZipArtifact(SourceCodeArtifact):
-    artifact_path: str
     main_file: str
+    qwak_artifact_path: Optional[str] = None
+    jfrog_artifact_path: Optional[JfrogArtifact] = None
+
+    @model_validator(mode="after")
+    def validate_only_one_artifact_path_type_provided(self: Self) -> Self:
+        """
+        Ensures that either 'qwak_artifact_path' or 'jfrog_artifact_path' is provided,
+        but not both.
+        """
+        if (self.qwak_artifact_path is None) == (self.jfrog_artifact_path is None):
+            raise QwakException(
+                "Either `qwak_artifact_path` or `jfrog_artifact_path` must be provided."
+            )
+        return self
 
     @classmethod
     def _from_proto(cls, proto: ProtoSourceCodeArtifact) -> "ZipArtifact":
         artifact_type: str = proto.WhichOneof("type")
         if artifact_type != "zip_artifact":
-            raise QwakException(f"Instead of `zip_artifact` got: {artifact_type}")
+            raise QwakException(
+                f"Instead of `zip_artifact` got: {artifact_type}"  # noqa
+            )
+
+        artifact_path_type: str = proto.zip_artifact.WhichOneof("path_type")
+        jfrog_artifact_path: Optional[JfrogArtifact] = None
+        qwak_artifact_path: Optional[str] = None
+
+        if artifact_path_type == "jfrog_artifact_path":
+            jfrog_artifact_path = JfrogArtifact.from_proto(
+                proto.zip_artifact.jfrog_artifact_path
+            )
+        elif artifact_path_type == "path":
+            qwak_artifact_path = proto.zip_artifact.path
+        else:
+            raise QwakException(
+                "Either `path` or `jfrog_artifact_path` must be provided inside SourceCodeArtifact proto."
+            )
 
         return cls(
-            artifact_path=proto.zip_artifact.path,
+            jfrog_artifact_path=jfrog_artifact_path,
+            qwak_artifact_path=qwak_artifact_path,
             main_file=proto.zip_artifact.main_file,
         )
 
     def _to_proto(self) -> ProtoSourceCodeArtifact:
         return ProtoSourceCodeArtifact(
             zip_artifact=ProtoZipArtifact(
-                path=self.artifact_path, main_file=self.main_file
+                jfrog_artifact_path=(
+                    self.jfrog_artifact_path.to_proto()
+                    if self.jfrog_artifact_path
+                    else None
+                ),
+                path=self.qwak_artifact_path,
+                main_file=self.main_file,
             )
         )
 
 
-@dataclass
-class SourceCodeSpec:
-    artifact: Optional[SourceCodeArtifact] = field(default=None)
+class SourceCodeSpec(BaseModel):
+    artifact: Optional[SourceCodeArtifact] = Field(default=None)
 
     @classmethod
     def _from_proto(cls, proto: ProtoSourceCodeSpec) -> "SourceCodeSpec":

qwak/feature_store/_common/source_code_spec_factory.py

@@ -1,23 +1,30 @@
 import os
 from pathlib import Path
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from qwak.exceptions import QwakException
 from qwak.feature_store._common import packaging
-from qwak.feature_store._common.source_code_spec import SourceCodeSpec, ZipArtifact
+from qwak.feature_store._common.source_code_spec import (
+    JfrogArtifact,
+    SourceCodeSpec,
+    ZipArtifact,
+)
 
 
 class SourceCodeSpecFactory:
     @staticmethod
-    def _upload_source_code_dir(main_entity_path: Path, presign_url: str):
+    def _upload_source_code_dir(
+        main_entity_path: Path, url: str, extra_headers: Dict[str, Any]
+    ):
         zip_location: Optional[Path] = None
         try:
             zip_location = packaging.zip_source_code_dir(base_path=main_entity_path)
             with open(zip_location, "rb") as zip_file:
-                packaging.put_presigned_url(
-                    presign_url=presign_url,
+                packaging.put_files_content(
+                    url=url,
                     content=zip_file.read(),
                     content_type="application/octet-stream",
+                    extra_headers=extra_headers,
                 )
         except Exception as e:
             raise QwakException("Got an error while trying to upload file.") from e
@@ -25,25 +32,71 @@ class SourceCodeSpecFactory:
             if zip_location:
                 os.remove(zip_location)
 
-    @staticmethod
+    @classmethod
     def get_zip_source_code_spec(
-        main_entity_path: Path, presign_url: str
+        cls,
+        main_entity_path: Path,
+        url: str,
+        extra_headers: Optional[Dict[str, str]] = None,
     ) -> SourceCodeSpec:
         """
         Zips and upload the artifact,
         the main "entity" relative path is set as the main_entity_path
         """
-        SourceCodeSpecFactory._upload_source_code_dir(
-            main_entity_path=main_entity_path, presign_url=presign_url
+        cls._upload_source_code_dir(
+            main_entity_path=main_entity_path,
+            url=url,
+            extra_headers=extra_headers or {},
         )
 
         parent_base_path = (
             main_entity_path.parent if main_entity_path.is_file() else main_entity_path
         )
 
+        if cls._is_artifactory_source_code(extra_headers):
+            return cls._build_artifactory_source_code_spec(
+                url=url,
+                parent_base_path=parent_base_path,
+                main_entity_path=main_entity_path,
+            )
+        else:
+            return cls._build_qwak_source_code_spec(
+                url=url,
+                parent_base_path=parent_base_path,
+                main_entity_path=main_entity_path,
+            )
+
+    @classmethod
+    def _is_artifactory_source_code(
+        cls, extra_headers: Optional[Dict[str, str]]
+    ) -> bool:
+        """
+        Check if the source code is from Artifactory by checking the extra headers.
+        """
+        if not extra_headers:
+            return False
+        return True
+
+    @classmethod
+    def _build_artifactory_source_code_spec(
+        cls, url: str, parent_base_path: Path, main_entity_path: Path
+    ) -> SourceCodeSpec:
+        return SourceCodeSpec(
+            artifact=ZipArtifact(
+                jfrog_artifact_path=JfrogArtifact(path=url),
+                main_file=str(
+                    Path(parent_base_path.name) / Path(main_entity_path.name)
+                ),
+            )
+        )
+
+    @classmethod
+    def _build_qwak_source_code_spec(
+        cls, url: str, parent_base_path: Path, main_entity_path: Path
+    ) -> SourceCodeSpec:
         return SourceCodeSpec(
             artifact=ZipArtifact(
-                artifact_path=presign_url,
+                qwak_artifact_path=url,
                 main_file=str(
                     Path(parent_base_path.name) / Path(main_entity_path.name)
                 ),

qwak/feature_store/data_sources/__init__.py

@@ -2,6 +2,7 @@ from qwak.feature_store.data_sources.batch.athena import AthenaSource
 from qwak.feature_store.data_sources.batch.big_query import BigQuerySource
 from qwak.feature_store.data_sources.batch.clickhouse import ClickhouseSource
 from qwak.feature_store.data_sources.batch.csv import CsvSource
+from qwak.feature_store.data_sources.batch.unity_catalog import UnityCatalogSource
 from qwak.feature_store.data_sources.batch.elastic_search import ElasticSearchSource
 from qwak.feature_store.data_sources.batch.filesystem.aws import (
     AnonymousS3Configuration,
@@ -52,6 +53,7 @@ __all__ = [
     "RedshiftSource",
     "SnowflakeSource",
     "VerticaSource",
+    "UnityCatalogSource",
     "KafkaSource",
     "Deserializer",
     "CustomDeserializer",

qwak/feature_store/data_sources/base.py

@@ -6,12 +6,13 @@ 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 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 (
@@ -63,14 +64,16 @@ class BaseSource(ABC):
             uploaded_artifact_url = self._upload_artifact()
 
         if source_definition_path:
-            presign_url: (
-                str
+            code_upload_args: (
+                GetDataSourceSourceCodeUploadResponse
             ) = FeatureRegistryClient().get_datasource_source_code_presign_url(
-                ds_name=self.name
+                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, presign_url=presign_url
+                    main_entity_path=source_definition_path,
+                    url=code_upload_args.upload_url,
+                    extra_headers=code_upload_args.extra_headers,
                 )
             )
 
@@ -113,6 +116,7 @@ class BaseSource(ABC):
             data_source=self,
             sample_size=number_of_rows,
             validation_options=validation_options,
+            silence_specific_exceptions=False,
         )
 
         if isinstance(response, SuccessValidationResponse):

qwak/feature_store/data_sources/batch/snowflake.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 from typing import Optional
+import warnings
 
 from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
     BatchSource as ProtoBatchSource,
@@ -11,15 +12,18 @@ from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
 from qwak.exceptions import QwakException
 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
@@ -28,16 +32,36 @@ class SnowflakeSource(BaseBatchSource):
         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!")
+
+        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(
@@ -50,6 +74,7 @@ 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,
@@ -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,
+        )