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

qwak/feature_store/_common/functions.py

@@ -1,10 +1,5 @@
 from typing import List
 
-from _qwak_proto.qwak.features_operator.v3.features_operator_pb2 import (
-    Outputs,
-    ValidationSuccessResponse,
-)
-
 
 def _generate_new_col(col_name: str, duplicate_cols_dict: dict):
     """
@@ -21,41 +16,3 @@ def _generate_new_col(col_name: str, duplicate_cols_dict: dict):
         return col_name
     elif len(feature_full_name) == 2 and duplicate_cols_dict[feature_full_name[1]] == 1:
         return feature_full_name[1]
-
-
-def normalize_cols(cols) -> List[str]:
-    """
-    Normalize cols - try to remove leading feature set name from features
-    Args:
-        cols: list of column
-    Return normalized columns
-    """
-    duplicate_columns_dict = {}
-    for col in cols:
-        feature_full_name: List[str] = col.split(".")
-        if len(feature_full_name) != 2:
-            duplicate_columns_dict[col] = 1
-        else:
-            fs_name = feature_full_name[1]
-            duplicate_columns_dict[fs_name] = duplicate_columns_dict.get(fs_name, 0) + 1
-    cols = [_generate_new_col(col_name, duplicate_columns_dict) for col_name in cols]
-    return cols
-
-
-def print_validation_outputs(response: ValidationSuccessResponse) -> bool:
-    did_print = False
-    outputs = response.outputs
-
-    if isinstance(outputs, Outputs) and (outputs.stdout or outputs.stderr):
-        message = "Validation outputs: "
-
-        if outputs.stdout:
-            message += f"stdout: {outputs.stdout}\n "
-
-        if outputs.stderr:
-            message += f"stderr: {outputs.stderr}"
-
-        print(message)
-        did_print = True
-
-    return did_print

qwak/feature_store/_common/packaging.py

@@ -0,0 +1,262 @@
+import contextlib
+import os
+import shutil
+import tempfile
+import zipfile
+from io import BytesIO
+from pathlib import Path
+from shutil import make_archive
+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"
+
+if TYPE_CHECKING:
+    from qwak.feature_store.data_sources.streaming.kafka.deserialization import (
+        Deserializer,
+    )
+    from qwak.feature_store.feature_sets.transformations.transformations import (
+        BaseTransformation,
+    )
+
+
+def zip_source_code_dir(base_path: Path) -> Path:
+    ignore_patterns = ["*.pyc", "__pycache__"]
+    return zip_tree_to_temp_dir(base_path=base_path, ignore_patterns=ignore_patterns)
+
+
+def zip_tree_to_temp_dir(base_path: Path, ignore_patterns=None) -> Path:
+    if ignore_patterns is None:
+        ignore_patterns = []
+
+    full_path_to_archive = None
+    parent_base_path = base_path.parent if base_path.is_file() else base_path
+    tempdir = tempfile.mkdtemp()
+    dst: Path = Path(os.path.join(tempdir, parent_base_path.name))
+    try:
+        shutil.copytree(
+            src=parent_base_path,
+            dst=dst,
+            ignore=shutil.ignore_patterns(*ignore_patterns),
+        )
+
+        full_path_to_archive: Path = zip_path_dir_recursively(
+            p=dst, dest=tempfile.mkdtemp()
+        )
+        return full_path_to_archive
+    except Exception as e:
+        raise QwakException("Failed to copy and zip directory") from e
+    finally:
+        shutil.rmtree(tempdir, ignore_errors=True)
+        if full_path_to_archive:
+            shutil.rmtree(full_path_to_archive, ignore_errors=True)
+
+
+def zip_path_dir_recursively(p: Path, dest: Path) -> Path:
+    root_dir: Path = p.parent if p.is_file() else p
+    full_path_to_archive: str = make_archive(
+        base_name=os.path.join(dest, str(uuid4())),
+        format="zip",
+        root_dir=root_dir.parent,
+        base_dir=root_dir.name,
+    )
+    return Path(full_path_to_archive)
+
+
+def zip_function_with_base_folder(
+    functions: List[Callable], fs_object_module_dir: Path = None
+) -> bytes:
+    """
+    Prepare zip bytes of the UDF functions and it's parent folder
+    In case the functions param is a Callable it will write the pkled code to "function.pkl" (original flow)
+    If a list of functions it will write each function to {function_name}.pkl
+    Args:
+        functions: udf functions
+        fs_object_module_dir: parent directory path of the registered feature store object
+    Returns:
+        bytes representing the zip bytes of the udf functions pkl and module
+    """
+    try:
+        import cloudpickle
+    except Exception:
+        raise QwakException(
+            f"Missing required 'cloudpickle' dependency. {UPDATE_QWAK_SDK_WITH_FEATURE_STORE_EXTRA_MSG}"
+        )
+
+    in_memory = BytesIO()
+    with chdir(fs_object_module_dir):
+        zf = zipfile.ZipFile(in_memory, mode="w")
+        with zf:
+            if fs_object_module_dir:
+                for path in Path.cwd().glob("**/*"):
+                    str_path = (
+                        f"./{path.relative_to(Path.cwd())}"
+                        if path.is_file()
+                        else str(path.resolve())
+                    )
+                    zf.write(str_path, "code" + "/" + str_path)
+            for function in functions:
+                zf.writestr(f"{function.__name__}.pkl", cloudpickle.dumps(function))
+        in_memory.seek(0)
+    return in_memory.read()
+
+
+@contextlib.contextmanager
+def chdir(dirname: Path = None):
+    """
+    chdir to given path and return to previous cwd when done
+    Args:
+        dirname: dir to chdir into
+    """
+    current_dir = Path.cwd()
+    try:
+        if dirname is not None:
+            os.chdir(dirname)
+        yield
+    finally:
+        os.chdir(current_dir)
+
+
+def get_featureset_presign_url(
+    features_manager_client: FeatureRegistryClient,
+    feature_store_object_name: str,
+    feature_store_object_name_suffix: str,
+    transformation,
+) -> str:
+    return features_manager_client.get_presigned_url_v1(
+        feature_set_name=feature_store_object_name + feature_store_object_name_suffix,
+        transformation=transformation._to_proto(),
+    )
+
+
+def get_data_source_presign_url_v1(
+    features_manager_client: FeatureRegistryClient,
+    data_source_object_name: str,
+    data_source_name_suffix: str,
+) -> str:
+    return features_manager_client.get_data_source_presigned_url_v1(
+        data_source_name=data_source_object_name, object_name=data_source_name_suffix
+    )
+
+
+def upload_artifact(
+    feature_store_object_name: str,
+    feature_store_object_name_suffix: str,
+    functions: List[Callable],
+    feature_module_dir: Path,
+    features_manager_client: FeatureRegistryClient,
+    artifact_object: Union["BaseTransformation", "Deserializer"],
+) -> str:
+    """
+    This function is used for uploading the function code in case a UDF type transformation was used (Pandas/Koalas/Pyspark)
+    If a transformation was supplied, meaning we're running for at least batch, will use the SaaS adapted
+    function to get the presign URL.
+    """
+    uploaded_artifact_url: str
+
+    # this is done here to avoid cyclic imports
+    from qwak.feature_store.data_sources.streaming.kafka.deserialization import (
+        Deserializer,
+    )
+    from qwak.feature_store.feature_sets.transformations.transformations import (
+        BaseTransformation,
+    )
+
+    if isinstance(artifact_object, BaseTransformation):
+        presign_url = get_featureset_presign_url(
+            features_manager_client=features_manager_client,
+            feature_store_object_name=feature_store_object_name,
+            feature_store_object_name_suffix=feature_store_object_name_suffix,
+            transformation=artifact_object,
+        )
+    elif isinstance(artifact_object, Deserializer):
+        presign_url = get_data_source_presign_url_v1(
+            features_manager_client=features_manager_client,
+            data_source_object_name=feature_store_object_name,
+            data_source_name_suffix=feature_store_object_name_suffix,
+        )
+    else:
+        raise QwakException(
+            f"Got an unsupported artifact type: {type(artifact_object)}"
+        )
+
+    uploaded_artifact_url = base_upload_function_code(
+        presign_url, functions, feature_module_dir
+    )
+
+    return uploaded_artifact_url
+
+
+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:
+    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,
+    )
+
+    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
+
+
+def base_upload_function_code(
+    presign_url: str,
+    functions: List[Callable],
+    feature_module_dir: Path,
+) -> str:
+    """
+    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
+        feature_module_dir: parent directory path of the feature
+    Returns:
+        Full path to the zip file uploaded
+    """
+    in_memory_zip = zip_function_with_base_folder(functions, feature_module_dir)
+    return upload_to_s3(
+        presign_url=presign_url,
+        in_memory_zip=in_memory_zip,
+        content_type=ZIP_FUNCTION_CONTENT_TYPE,
+    )

qwak/feature_store/_common/source_code_spec.py

@@ -0,0 +1,108 @@
+from abc import ABC, abstractmethod
+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 (
+    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(BaseModel, ABC):
+    @abstractmethod
+    def _to_proto(self) -> ProtoSourceCodeArtifact:
+        pass
+
+
+class ZipArtifact(SourceCodeArtifact):
+    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}"  # 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(
+            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(
+                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,
+            )
+        )
+
+
+class SourceCodeSpec(BaseModel):
+    artifact: Optional[SourceCodeArtifact] = Field(default=None)
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoSourceCodeSpec) -> "SourceCodeSpec":
+        artifact_type: str = proto.artifact.WhichOneof("type")
+        if artifact_type == "zip_artifact":
+            return cls(artifact=ZipArtifact._from_proto(proto.artifact))
+
+        return cls(artifact=None)
+
+    def _to_proto(self) -> ProtoSourceCodeSpec:
+        if self.artifact:
+            return ProtoSourceCodeSpec(artifact=self.artifact._to_proto())
+        return ProtoSourceCodeSpec()

qwak/feature_store/_common/source_code_spec_factory.py

@@ -0,0 +1,104 @@
+import os
+from pathlib import Path
+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 (
+    JfrogArtifact,
+    SourceCodeSpec,
+    ZipArtifact,
+)
+
+
+class SourceCodeSpecFactory:
+    @staticmethod
+    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_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
+        finally:
+            if zip_location:
+                os.remove(zip_location)
+
+    @classmethod
+    def get_zip_source_code_spec(
+        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
+        """
+        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(
+                qwak_artifact_path=url,
+                main_file=str(
+                    Path(parent_base_path.name) / Path(main_entity_path.name)
+                ),
+            )
+        )

qwak/feature_store/_common/value.py

@@ -0,0 +1,10 @@
+UPDATE_QWAK_SDK_WITH_FEATURE_STORE_EXTRA_MSG = (
+    'please run: pip install -U "qwak-sdk[feature-store]"'
+)
+
+
+def missing_pyspark_exception_message(name: str) -> str:
+    return (
+        f"Missing 'pyspark' dependency required for {name} transformation. "
+        f"{UPDATE_QWAK_SDK_WITH_FEATURE_STORE_EXTRA_MSG}"
+    )

qwak/feature_store/data_sources/__init__.py

@@ -1,31 +1,67 @@
-from qwak.feature_store.data_sources.batch_sources.big_query import BigQuerySource
-from qwak.feature_store.data_sources.batch_sources.csv import CsvSource
-from qwak.feature_store.data_sources.batch_sources.elastic_search import (
-    ElasticSearchSource,
-)
-from qwak.feature_store.data_sources.batch_sources.filesystem_config import (
+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,
+    AwsS3AssumeRoleFileSystemConfiguration,
     AwsS3FileSystemConfiguration,
 )
-from qwak.feature_store.data_sources.batch_sources.mongodb import MongoDbSource
-from qwak.feature_store.data_sources.batch_sources.mysql import MysqlSource
-from qwak.feature_store.data_sources.batch_sources.parquet import ParquetSource
-from qwak.feature_store.data_sources.batch_sources.postgres import PostgresSource
-from qwak.feature_store.data_sources.batch_sources.redshift import RedshiftSource
-from qwak.feature_store.data_sources.batch_sources.snowflake import SnowflakeSource
-from qwak.feature_store.data_sources.batch_sources.vertica import VerticaSource
+from qwak.feature_store.data_sources.batch.filesystem.gcp import (
+    GcpGcsServiceAccountImpersonation,
+    GcpGcsUnauthenticated,
+)
+from qwak.feature_store.data_sources.batch.mongodb import MongoDbSource
+from qwak.feature_store.data_sources.batch.mysql import MysqlSource
+from qwak.feature_store.data_sources.batch.parquet import ParquetSource
+from qwak.feature_store.data_sources.batch.postgres import PostgresSource
+from qwak.feature_store.data_sources.batch.redshift import RedshiftSource
+from qwak.feature_store.data_sources.batch.snowflake import SnowflakeSource
+from qwak.feature_store.data_sources.batch.vertica import VerticaSource
+from qwak.feature_store.data_sources.streaming.kafka import KafkaSource
+from qwak.feature_store.data_sources.streaming.kafka.authentication import (
+    PlainAuthentication,
+    SaslAuthentication,
+    SaslMechanism,
+    SecurityProtocol,
+    SslAuthentication,
+)
+from qwak.feature_store.data_sources.streaming.kafka.deserialization import (
+    CustomDeserializer,
+    Deserializer,
+    GenericDeserializer,
+    MessageFormat,
+)
 
 __all__ = [
+    "AthenaSource",
     "BigQuerySource",
+    "ClickhouseSource",
     "CsvSource",
+    "ElasticSearchSource",
+    "AwsS3FileSystemConfiguration",
+    "AnonymousS3Configuration",
+    "AwsS3AssumeRoleFileSystemConfiguration",
+    "GcpGcsServiceAccountImpersonation",
+    "GcpGcsUnauthenticated",
     "MongoDbSource",
     "MysqlSource",
     "ParquetSource",
-    "RedshiftSource",
     "PostgresSource",
+    "RedshiftSource",
     "SnowflakeSource",
     "VerticaSource",
-    "ElasticSearchSource",
-    "AwsS3FileSystemConfiguration",
-    "AnonymousS3Configuration",
+    "UnityCatalogSource",
+    "KafkaSource",
+    "Deserializer",
+    "CustomDeserializer",
+    "GenericDeserializer",
+    "MessageFormat",
+    "PlainAuthentication",
+    "SslAuthentication",
+    "SaslMechanism",
+    "SaslAuthentication",
+    "SecurityProtocol",
 ]

qwak/feature_store/data_sources/attributes.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass, field
+
+from _qwak_proto.qwak.feature_store.sources.data_source_attribute_pb2 import (
+    DataSourceAttributes as ProtoDataSourceAttributes,
+    DataSourceAttributes,
+)
+from qwak.feature_store._common.source_code_spec import SourceCodeSpec
+
+
+@dataclass
+class DataSourceAttributes:
+    source_code_spec: SourceCodeSpec = field(default_factory=lambda: SourceCodeSpec())
+
+    def _to_proto(self) -> ProtoDataSourceAttributes:
+        return ProtoDataSourceAttributes(
+            source_code_spec=self.source_code_spec._to_proto()
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoDataSourceAttributes) -> DataSourceAttributes:
+        return cls(
+            source_code_spec=SourceCodeSpec._from_proto(proto=proto.source_code_spec)
+        )