frogml-core 0.0.1__py3-none-any.whl

frogml/feature_store/_common/packaging.py

@@ -0,0 +1,244 @@
+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, List, Union
+from uuid import uuid4
+
+import requests
+
+from frogml.clients.feature_store.management_client import FeatureRegistryClient
+from frogml.exceptions import FrogmlException
+from frogml.feature_store._common.value import (
+    UPDATE_FROGML_SDK_WITH_FEATURE_STORE_EXTRA_MSG,
+)
+
+ZIP_FUNCTION_CONTENT_TYPE = "text/plain"
+
+if TYPE_CHECKING:
+    from frogml.feature_store.data_sources.streaming.kafka.deserialization import (
+        Deserializer,
+    )
+    from frogml.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 FrogmlException("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 FrogmlException(
+            f"Missing required 'cloudpickle' dependency. {UPDATE_FROGML_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 frogml.feature_store.data_sources.streaming.kafka.deserialization import (
+        Deserializer,
+    )
+    from frogml.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 FrogmlException(
+            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_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 FrogmlException(
+            f"Error while uploading user function to s3. status code: {http_response.status_code}"
+        )
+
+
+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
+    )
+    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
+    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 and upload 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,
+    )

frogml/feature_store/_common/source_code_spec.py

@@ -0,0 +1,62 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Optional
+
+from _frogml_proto.qwak.feature_store.v1.common.source_code.source_code_pb2 import (
+    SourceCodeArtifact as ProtoSourceCodeArtifact,
+)
+from _frogml_proto.qwak.feature_store.v1.common.source_code.source_code_pb2 import (
+    SourceCodeSpec as ProtoSourceCodeSpec,
+)
+from _frogml_proto.qwak.feature_store.v1.common.source_code.source_code_pb2 import (
+    ZipArtifact as ProtoZipArtifact,
+)
+from frogml.exceptions import FrogmlException
+
+
+class SourceCodeArtifact(ABC):
+    @abstractmethod
+    def _to_proto(self) -> ProtoSourceCodeArtifact:
+        pass
+
+
+@dataclass
+class ZipArtifact(SourceCodeArtifact):
+    artifact_path: str
+    main_file: str
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoSourceCodeArtifact) -> "ZipArtifact":
+        artifact_type: str = proto.WhichOneof("type")
+        if artifact_type != "zip_artifact":
+            raise FrogmlException(f"Instead of `zip_artifact` got: {artifact_type}")
+
+        return cls(
+            artifact_path=proto.zip_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
+            )
+        )
+
+
+@dataclass
+class SourceCodeSpec:
+    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()

frogml/feature_store/_common/source_code_spec_factory.py

@@ -0,0 +1,51 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from frogml.exceptions import FrogmlException
+from frogml.feature_store._common import packaging
+from frogml.feature_store._common.source_code_spec import SourceCodeSpec, ZipArtifact
+
+
+class SourceCodeSpecFactory:
+    @staticmethod
+    def _upload_source_code_dir(main_entity_path: Path, presign_url: str):
+        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,
+                    content=zip_file.read(),
+                    content_type="application/octet-stream",
+                )
+        except Exception as e:
+            raise FrogmlException("Got an error while trying to upload file.") from e
+        finally:
+            if zip_location:
+                os.remove(zip_location)
+
+    @staticmethod
+    def get_zip_source_code_spec(
+        main_entity_path: Path, presign_url: str
+    ) -> 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
+        )
+
+        parent_base_path = (
+            main_entity_path.parent if main_entity_path.is_file() else main_entity_path
+        )
+
+        return SourceCodeSpec(
+            artifact=ZipArtifact(
+                artifact_path=presign_url,
+                main_file=str(
+                    Path(parent_base_path.name) / Path(main_entity_path.name)
+                ),
+            )
+        )

frogml/feature_store/_common/value.py

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

frogml/feature_store/data_sources/__init__.py

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

frogml/feature_store/data_sources/attributes.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass, field
+
+from _frogml_proto.qwak.feature_store.sources.data_source_attribute_pb2 import (
+    DataSourceAttributes,
+)
+from _frogml_proto.qwak.feature_store.sources.data_source_attribute_pb2 import (
+    DataSourceAttributes as ProtoDataSourceAttributes,
+)
+from frogml.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)
+        )

frogml/feature_store/data_sources/base.py

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

frogml/feature_store/data_sources/batch/_batch.py

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

frogml/feature_store/data_sources/batch/_jdbc.py

@@ -0,0 +1,22 @@
+from abc import ABC
+from dataclasses import dataclass
+from typing import Optional
+
+from frogml.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
+
+    def __post_init__(self):
+        self._validate()
+
+    def _validate(self):
+        if not (bool(self.db_table) ^ bool(self.query)):
+            raise FrogmlException("Only one of query and db_table must be set")