frogml-core 0.0.116__py3-none-any.whl → 0.0.117__py3-none-any.whl

frogml_core/feature_store/_common/packaging.py

@@ -1,244 +0,0 @@
-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_core.clients.feature_store.management_client import FeatureRegistryClient
-from frogml_core.exceptions import FrogmlException
-from frogml_core.feature_store._common.value import (
-    UPDATE_FROGML_SDK_WITH_FEATURE_STORE_EXTRA_MSG,
-)
-
-ZIP_FUNCTION_CONTENT_TYPE = "text/plain"
-
-if TYPE_CHECKING:
-    from frogml_core.feature_store.data_sources.streaming.kafka.deserialization import (
-        Deserializer,
-    )
-    from frogml_core.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_core.feature_store.data_sources.streaming.kafka.deserialization import (
-        Deserializer,
-    )
-    from frogml_core.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_core/feature_store/_common/source_code_spec.py

@@ -1,62 +0,0 @@
-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_core.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_core/feature_store/_common/source_code_spec_factory.py

@@ -1,54 +0,0 @@
-import os
-from pathlib import Path
-from typing import Optional
-
-from frogml_core.exceptions import FrogmlException
-from frogml_core.feature_store._common import packaging
-from frogml_core.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_core/feature_store/data_sources/__init__.py

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

@@ -1,25 +0,0 @@
-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_core.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_core/feature_store/data_sources/base.py

@@ -1,128 +0,0 @@
-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_core.clients.feature_store import FeatureRegistryClient
-from frogml_core.exceptions import FrogmlException
-from frogml_core.feature_store._common.artifact_utils import (
-    ArtifactSpec,
-    ArtifactsUploader,
-)
-from frogml_core.feature_store._common.source_code_spec import SourceCodeSpec
-from frogml_core.feature_store._common.source_code_spec_factory import (
-    SourceCodeSpecFactory,
-)
-from frogml_core.feature_store.data_sources.attributes import DataSourceAttributes
-from frogml_core.feature_store.validations.validation_options import (
-    DataSourceValidationOptions,
-)
-from frogml_core.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_core.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_core/feature_store/data_sources/batch/_batch.py

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

frogml_core/feature_store/data_sources/batch/_jdbc.py

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