qwak-core 0.4.378__py3-none-any.whl → 0.5.12__py3-none-any.whl

qwak/feature_store/feature_sets/streaming_backfill.py

@@ -1,24 +1,18 @@
-from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from typing import List, Optional, Set, Union
 
 from _qwak_proto.qwak.feature_store.features.execution_pb2 import (
-    BackfillExecutionSpec as ProtoBackfillExecutionSpec,
+    ExecutionSpec as ProtoExecutionSpec,
 )
-from _qwak_proto.qwak.feature_store.features.feature_set_types_pb2 import (
-    BackfillBatchDataSourceSpec as ProtoBackfillBatchDataSourceSpec,
-    BackfillDataSourceSpec as ProtoBackfillDataSourceSpec,
-    BackfillSpec as ProtoBackfillSpec,
-)
-from _qwak_proto.qwak.feature_store.sources.batch_pb2 import (
-    BatchSource as ProtoBatchSource,
+from _qwak_proto.qwak.execution.v1.streaming_aggregation_pb2 import (
+    StreamingAggregationBackfillIngestion as ProtoStreamingAggregationBackfillIngestion,
+    BackfillDataSource as ProtoBackfillDataSource,
+    TimeRange as ProtoTimeRange,
 )
 from google.protobuf.timestamp_pb2 import Timestamp as ProtoTimestamp
-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.feature_set_utils import get_batch_source_for_featureset
 from qwak.feature_store.feature_sets.execution_spec import ClusterTemplate
 from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
 
@@ -26,36 +20,15 @@ _BACKFILL_ = "_qwak_backfill_specification"
 
 
 @dataclass
-class DataSourceBackfillSpec(ABC):
+class BackfillDataSource:
     data_source_name: str
-
-    @abstractmethod
-    def _to_proto(self, feature_registry: FeatureRegistryClient):
-        pass
-
-    @classmethod
-    def _from_proto(cls, proto: ProtoBackfillDataSourceSpec):
-        function_mapping = {"batch_data_source_spec": BackfillBatchDataSourceSpec}
-
-        backfill_source_type: str = proto.WhichOneof("type")
-
-        if backfill_source_type in function_mapping:
-            function_class = function_mapping.get(backfill_source_type)
-            return function_class._from_proto(proto)
-
-        raise QwakException(
-            f"Got unsupported backfill source type {backfill_source_type} for streaming backfill"
-        )
-
-
-@dataclass
-class BackfillBatchDataSourceSpec(DataSourceBackfillSpec):
     start_datetime: Optional[datetime] = None
     end_datetime: Optional[datetime] = None
 
-    def _to_proto(
-        self, feature_registry: FeatureRegistryClient
-    ) -> ProtoBackfillBatchDataSourceSpec:
+    def __post_init__(self):
+        self._validate()
+
+    def _to_proto(self) -> ProtoBackfillDataSource:
         start_timestamp: Optional[ProtoTimestamp] = None
         end_timestamp: Optional[ProtoTimestamp] = None
 
@@ -67,63 +40,94 @@ class BackfillBatchDataSourceSpec(DataSourceBackfillSpec):
             start_timestamp = ProtoTimestamp()
             start_timestamp.FromDatetime(self.start_datetime.astimezone(timezone.utc))
 
-        proto_data_source: ProtoBatchSource = get_batch_source_for_featureset(
-            batch_ds_name=self.data_source_name, feature_registry=feature_registry
-        )
-
-        return ProtoBackfillBatchDataSourceSpec(
-            data_source=proto_data_source,
+        time_range = ProtoTimeRange(
             start_timestamp=start_timestamp,
             end_timestamp=end_timestamp,
         )
 
+        return ProtoBackfillDataSource(
+            data_source_name=self.data_source_name,
+            time_range=time_range,
+        )
+
     @classmethod
-    def _from_proto(
-        cls, proto: ProtoBackfillDataSourceSpec
-    ) -> "BackfillBatchDataSourceSpec":
+    def _from_proto(cls, proto: ProtoBackfillDataSource) -> "BackfillDataSource":
         start_datetime: Optional[datetime] = None
         end_datetime: Optional[datetime] = None
 
-        batch_backfill_spec: ProtoBackfillBatchDataSourceSpec = (
-            proto.batch_data_source_spec
-        )
+        time_range: ProtoTimeRange = proto.time_range
 
-        proto_start_timestamp: ProtoTimestamp = batch_backfill_spec.start_timestamp
-        proto_end_timestamp: ProtoTimestamp = batch_backfill_spec.end_timestamp
+        proto_start_timestamp: Optional[ProtoTimestamp] = (
+            time_range.start_timestamp if time_range.start_timestamp else None
+        )
+        proto_end_timestamp: Optional[ProtoTimestamp] = (
+            time_range.end_timestamp if time_range.end_timestamp else None
+        )
 
-        start_datetime = datetime.fromtimestamp(
-            proto_start_timestamp.seconds + proto_start_timestamp.nanos / 1e9
+        start_datetime = (
+            datetime.fromtimestamp(
+                proto_start_timestamp.seconds + proto_start_timestamp.nanos / 1e9
+            )
+            if proto_start_timestamp
+            else None
         )
 
-        end_datetime = datetime.fromtimestamp(
-            proto_end_timestamp.seconds + proto_end_timestamp.nanos / 1e9
+        end_datetime = (
+            datetime.fromtimestamp(
+                proto_end_timestamp.seconds + proto_end_timestamp.nanos / 1e9
+            )
+            if proto_end_timestamp
+            else None
         )
 
         return cls(
-            data_source_name=batch_backfill_spec.data_source.name,
+            data_source_name=proto.data_source_name,
             start_datetime=start_datetime,
             end_datetime=end_datetime,
         )
 
+    def _validate(self):
+        if self.start_datetime and self.end_datetime:
+            if self.start_datetime >= self.end_datetime:
+                raise QwakException(
+                    f"Backfill data source {self.data_source_name} has invalid time range: "
+                    f"start_datetime {self.start_datetime} is after or equal end_datetime {self.end_datetime}."
+                )
+
+        if not self.data_source_name:
+            raise QwakException(
+                "Backfill data source must have a valid data source name."
+            )
+
 
 @dataclass
 class StreamingBackfill:
+    featureset_name: str
     start_datetime: datetime
     end_datetime: datetime
-    data_sources_specs: List[DataSourceBackfillSpec]
+    data_sources: List[BackfillDataSource]
     transform: "SparkSqlTransformation"
     cluster_template: Optional[ClusterTemplate] = ClusterTemplate.SMALL
 
     def __post_init__(self):
-        if not self.data_sources_specs:
+        if not self.featureset_name:
+            raise QwakException("featureset_name must be provided for backfill.")
+
+        if not self.start_datetime or not self.end_datetime:
             raise QwakException(
-                "Trying to create a streaming backfill with no data sources. "
-                "At least one data source has to be provided when trying to create a streaming backfill."
+                "For Streaming backfill, start_datetime and end_datetime are mandatory fields."
             )
 
-        if not self.start_datetime or not self.end_datetime:
+        if self.start_datetime >= self.end_datetime:
             raise QwakException(
-                "For backfill, start_datetime and end_datetime are mandatory fields."
+                f"Backfill has invalid time range: "
+                f"start_datetime {self.start_datetime} is after or equal end_datetime {self.end_datetime}."
+            )
+
+        if not self.data_sources:
+            raise QwakException(
+                "Trying to create a streaming backfill with no data sources. "
+                "At least one data source has to be provided when trying to create a streaming backfill."
             )
 
         if type(self.transform) is not SparkSqlTransformation:
@@ -135,7 +139,7 @@ class StreamingBackfill:
 
     def _validate_unique_sources(self):
         source_names: List[str] = [
-            data_source.data_source_name for data_source in self.data_sources_specs
+            data_source.data_source_name for data_source in self.data_sources
         ]
         duplicates: Set[str] = {
             item for item in source_names if source_names.count(item) > 1
@@ -146,23 +150,14 @@ class StreamingBackfill:
                 f"Found these duplicates: {', '.join(set(duplicates))}"
             )
 
-    def _validate_tile_size(self, initial_tile_size: int):
-        if self.end_datetime.timestamp() % initial_tile_size != 0:
-            raise QwakException(
-                f"Chosen backfill end datetime is invalid,"
-                f" it has to be exactly dividable by slice size of {initial_tile_size} seconds."
-            )
-
     def _to_proto(
         self,
-        feature_registry: FeatureRegistryClient,
-        featureset_name: str,
         original_instance_module_path: str,
-    ) -> ProtoBackfillSpec:
+    ) -> ProtoStreamingAggregationBackfillIngestion:
         artifact_url: Optional[str] = None
         artifact_spec: Optional[ArtifactSpec] = ArtifactsUploader.get_artifact_spec(
             transformation=self.transform,
-            featureset_name=f"{featureset_name}-backfill",
+            featureset_name=f"{self.featureset_name}-backfill",
             __instance_module_path__=original_instance_module_path,
         )
 
@@ -175,85 +170,54 @@ class StreamingBackfill:
         start_timestamp = ProtoTimestamp()
         start_timestamp.FromDatetime(self.start_datetime.astimezone(timezone.utc))
 
-        return ProtoBackfillSpec(
+        return ProtoStreamingAggregationBackfillIngestion(
+            featureset_name=self.featureset_name,
             start_timestamp=start_timestamp,
             end_timestamp=end_timestamp,
-            execution_spec=ProtoBackfillExecutionSpec(
-                **{"cluster_template": ClusterTemplate.to_proto(self.cluster_template)}
+            execution_spec=ProtoExecutionSpec(
+                cluster_template=ClusterTemplate.to_proto(self.cluster_template)
             ),
             transformation=self.transform._to_proto(artifact_path=artifact_url),
             data_source_specs=[
-                ProtoBackfillDataSourceSpec(
-                    batch_data_source_spec=data_source_spec._to_proto(
-                        feature_registry=feature_registry
-                    )
-                )
-                for data_source_spec in self.data_sources_specs
+                data_source._to_proto() for data_source in self.data_sources
             ],
         )
 
     @classmethod
-    def _from_proto(cls, proto: ProtoBackfillSpec):
-        datetime.fromtimestamp(
-            proto.start_timestamp.seconds + proto.start_timestamp.nanos / 1e9
-        )
-
-        data_sources_specs = [
-            BackfillBatchDataSourceSpec._from_proto(ds)
-            for ds in proto.data_source_specs
+    def _from_proto(cls, proto: ProtoStreamingAggregationBackfillIngestion):
+        backfill_data_sources = [
+            BackfillDataSource._from_proto(ds) for ds in proto.data_source_specs
         ]
 
         return cls(
+            featureset_name=proto.featureset_name,
             start_datetime=datetime.fromtimestamp(
                 proto.start_timestamp.seconds + proto.start_timestamp.nanos / 1e9
             ),
             end_datetime=datetime.fromtimestamp(
                 proto.end_timestamp.seconds + proto.end_timestamp.nanos / 1e9
             ),
-            data_sources_specs=data_sources_specs,
+            data_sources=backfill_data_sources,
             transform=SparkSqlTransformation._from_proto(
                 proto.transformation.sql_transformation
             ),
+            cluster_template=(
+                ClusterTemplate.from_proto(proto.execution_spec.cluster_template)
+                if proto.execution_spec.cluster_template
+                else None
+            ),
         )
 
     @staticmethod
     def _get_normalized_backfill_sources_spec(
-        data_sources: Union[List[str], List[DataSourceBackfillSpec]],
-    ) -> List[DataSourceBackfillSpec]:
-        # reformat all data source specs to 'DataSourceBackfillSpec'
+        data_sources: Union[List[str], List[BackfillDataSource]],
+    ) -> List[BackfillDataSource]:
+        # reformat all data source names to 'BackfillDataSource'
         return [
             (
-                BackfillBatchDataSourceSpec(data_source_name=data_source)
+                BackfillDataSource(data_source_name=data_source)
                 if isinstance(data_source, str)
                 else data_source
             )
             for data_source in data_sources
         ]
-
-    @classmethod
-    def set_streaming_backfill_on_function(
-        cls,
-        function,
-        start_date: datetime,
-        end_date: datetime,
-        data_sources: Union[List[str], List[DataSourceBackfillSpec]],
-        backfill_transformation: SparkSqlTransformation,
-        backfill_cluster_template: Optional[ClusterTemplate] = ClusterTemplate.SMALL,
-    ):
-        setattr(
-            function,
-            _BACKFILL_,
-            cls(
-                start_datetime=start_date,
-                end_datetime=end_date,
-                data_sources_specs=StreamingBackfill._get_normalized_backfill_sources_spec(
-                    data_sources
-                ),
-                transform=backfill_transformation,
-                cluster_template=backfill_cluster_template,
-            ),
-        )
-
-    @staticmethod
-    def get_streaming_backfill_from_function(function):
-        return getattr(function, _BACKFILL_, None)

qwak/inner/const.py

@@ -33,10 +33,6 @@ class QwakConstants:
 
     TOKEN_AUDIENCE: str = "https://auth-token.qwak.ai/"  # nosec B105
 
-    QWAK_AUTHENTICATION_URL = "https://grpc.qwak.ai/api/v1/authentication/qwak-api-key"
-
-    QWAK_AUTHENTICATED_USER_ENDPOINT: str = (
-        "https://grpc.qwak.ai/api/v0/runtime/get-authenticated-user-context"
-    )
-
     QWAK_APP_URL: str = "https://app.qwak.ai"
+
+    CONTROL_PLANE_GRPC_ADDRESS_ENVAR_NAME: str = "CONTROL_PLANE_GRPC_ADDRESS"

qwak/inner/di_configuration/__init__.py

@@ -1,69 +1,3 @@
-import os
-
 from .account import UserAccountConfiguration
 from .containers import QwakContainer
-
-
-def wire_dependencies():
-    container = QwakContainer()
-
-    default_config_file = os.path.join(os.path.dirname(__file__), "config.yml")
-    container.config.from_yaml(default_config_file)
-
-    from qwak.clients import (
-        administration,
-        alert_management,
-        alerts_registry,
-        analytics,
-        audience,
-        automation_management,
-        autoscaling,
-        batch_job_management,
-        build_orchestrator,
-        data_versioning,
-        deployment,
-        feature_store,
-        file_versioning,
-        instance_template,
-        integration_management,
-        kube_deployment_captain,
-        logging_client,
-        model_management,
-        project,
-        prompt_manager,
-        system_secret,
-        user_application_instance,
-        vector_store,
-        workspace_manager,
-    )
-
-    container.wire(
-        packages=[
-            administration,
-            alert_management,
-            audience,
-            automation_management,
-            autoscaling,
-            analytics,
-            batch_job_management,
-            build_orchestrator,
-            data_versioning,
-            deployment,
-            file_versioning,
-            instance_template,
-            kube_deployment_captain,
-            logging_client,
-            model_management,
-            project,
-            feature_store,
-            user_application_instance,
-            alerts_registry,
-            workspace_manager,
-            vector_store,
-            integration_management,
-            system_secret,
-            prompt_manager,
-        ]
-    )
-
-    return container
+from .dependency_wiring import wire_dependencies

qwak/inner/di_configuration/dependency_wiring.py

@@ -0,0 +1,98 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from qwak.inner.const import QwakConstants
+from qwak.inner.di_configuration import QwakContainer
+from qwak.inner.tool.grpc.grpc_tools import validate_grpc_address
+from qwak.tools.logger import get_qwak_logger
+
+
+logger = get_qwak_logger()
+
+__DEFAULT_CONFIG_FILE_PATH: Path = Path(__file__).parent / "config.yml"
+
+
+def wire_dependencies():
+    container = QwakContainer()
+
+    container.config.from_yaml(__DEFAULT_CONFIG_FILE_PATH)
+    control_plane_grpc_address_override: Optional[str] = os.getenv(
+        QwakConstants.CONTROL_PLANE_GRPC_ADDRESS_ENVAR_NAME
+    )
+
+    if control_plane_grpc_address_override:
+        validate_grpc_address(control_plane_grpc_address_override)
+        __override_control_plane_grpc_address(
+            container, control_plane_grpc_address_override
+        )
+
+    from qwak.clients import (
+        administration,
+        alert_management,
+        alerts_registry,
+        analytics,
+        audience,
+        automation_management,
+        autoscaling,
+        batch_job_management,
+        build_orchestrator,
+        data_versioning,
+        deployment,
+        feature_store,
+        file_versioning,
+        instance_template,
+        integration_management,
+        kube_deployment_captain,
+        logging_client,
+        model_management,
+        project,
+        prompt_manager,
+        system_secret,
+        user_application_instance,
+        vector_store,
+        workspace_manager,
+    )
+
+    container.wire(
+        packages=[
+            administration,
+            alert_management,
+            audience,
+            automation_management,
+            autoscaling,
+            analytics,
+            batch_job_management,
+            build_orchestrator,
+            data_versioning,
+            deployment,
+            file_versioning,
+            instance_template,
+            kube_deployment_captain,
+            logging_client,
+            model_management,
+            project,
+            feature_store,
+            user_application_instance,
+            alerts_registry,
+            workspace_manager,
+            vector_store,
+            integration_management,
+            system_secret,
+            prompt_manager,
+        ]
+    )
+
+    return container
+
+
+def __override_control_plane_grpc_address(
+    container: "QwakContainer", control_plane_grpc_address_override: str
+):
+    logger.debug(
+        "Overriding control plane gRPC address from environment variable to %s.",
+        control_plane_grpc_address_override,
+    )
+    container.config.grpc.core.address.from_value(
+        control_plane_grpc_address_override.strip()
+    )

qwak/inner/tool/grpc/grpc_tools.py

@@ -1,15 +1,18 @@
 import logging
+import re
 import time
 from abc import ABC, abstractmethod
 from random import randint
 from typing import Callable, Optional, Tuple
+from urllib.parse import urlparse, ParseResult
 
 import grpc
-from qwak.exceptions import QwakException
 
+from qwak.exceptions import QwakException, QwakGrpcAddressException
 from .grpc_auth import Auth0Client
 
 logger = logging.getLogger()
+HOSTNAME_REGEX: str = r"^(?!-)(?:[A-Za-z0-9-]{1,63}\.)*[A-Za-z0-9-]{1,63}(?<!-)$"
 
 
 def create_grpc_channel(
@@ -19,7 +22,7 @@ def create_grpc_channel(
     auth_metadata_plugin: grpc.AuthMetadataPlugin = None,
     timeout: int = 100,
     options=None,
-    backoff_options={},
+    backoff_options=None,
     max_attempts=4,
     status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
     attempt=0,
@@ -40,6 +43,9 @@ def create_grpc_channel(
         status_for_retry: grpc statuses to retry upon
     Returns: Returns a grpc.Channel
     """
+    if backoff_options is None:
+        backoff_options = {}
+
     if not url:
         raise QwakException("Unable to create gRPC channel. URL has not been defined.")
 
@@ -101,11 +107,14 @@ def create_grpc_channel_or_none(
     auth_metadata_plugin: grpc.AuthMetadataPlugin = None,
     timeout: int = 30,
     options=None,
-    backoff_options={},
+    backoff_options=None,
     max_attempts=2,
     status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
     attempt=0,
 ) -> Callable[[Optional[str], Optional[bool]], Optional[grpc.Channel]]:
+    if backoff_options is None:
+        backoff_options = {}
+
     def deferred_channel(
         url_overwrite: Optional[str] = None, ssl_overwrite: Optional[bool] = None
     ):
@@ -129,6 +138,117 @@ def create_grpc_channel_or_none(
     return deferred_channel
 
 
+def validate_grpc_address(
+    grpc_address: str,
+    is_port_specification_allowed: bool = False,
+    is_url_scheme_allowed: bool = False,
+):
+    """
+    Validate gRPC address format
+    Args:
+        grpc_address (str): gRPC address to validate
+        is_port_specification_allowed (bool): Whether to allow port specification in the address
+        is_url_scheme_allowed (bool): Whether to allow URL scheme in the address
+    Raises:
+        QwakGrpcAddressException: If the gRPC address is invalid
+    """
+    parsed_grpc_address: ParseResult = parse_address(grpc_address)
+    hostname: str = get_hostname_from_address(parsed_grpc_address)
+    validate_paths_are_not_included_in_address(parsed_grpc_address)
+
+    if not is_url_scheme_allowed:
+        __validate_url_scheme_not_included_in_address(parsed_grpc_address)
+
+    if not is_port_specification_allowed:
+        __validate_port_not_included_in_address(parsed_grpc_address)
+
+    if not is_valid_hostname(hostname):
+        raise QwakGrpcAddressException(
+            "gRPC address must be a simple hostname or fully qualified domain name.",
+            parsed_grpc_address,
+        )
+
+
+def validate_paths_are_not_included_in_address(
+    parsed_grpc_address: ParseResult,
+) -> None:
+    has_invalid_path: bool = (
+        parsed_grpc_address.path not in {"", "/"}
+        or parsed_grpc_address.query
+        or parsed_grpc_address.fragment
+    )
+
+    if has_invalid_path:
+        raise QwakGrpcAddressException(
+            "gRPC address must not contain paths, queries, or fragments.",
+            parsed_grpc_address,
+        )
+
+
+def get_hostname_from_address(parsed_grpc_address: ParseResult) -> str:
+    hostname: Optional[str] = parsed_grpc_address.hostname
+    if not hostname:
+        raise QwakGrpcAddressException(
+            "gRPC address must contain a valid hostname.", parsed_grpc_address
+        )
+
+    return hostname
+
+
+def __validate_url_scheme_not_included_in_address(
+    parsed_grpc_address: ParseResult,
+) -> None:
+    if parsed_grpc_address.scheme:
+        raise QwakGrpcAddressException(
+            "URL scheme is not allowed in the gRPC address.", parsed_grpc_address
+        )
+
+
+def __validate_port_not_included_in_address(parsed_grpc_address: ParseResult):
+    try:
+        port: Optional[int] = parsed_grpc_address.port
+    except ValueError as exc:
+        raise QwakGrpcAddressException(
+            "Invalid port specification in the gRPC address.", parsed_grpc_address
+        ) from exc
+
+    if port:
+        raise QwakGrpcAddressException(
+            "Port specification is not allowed in the gRPC address.",
+            parsed_grpc_address,
+        )
+
+
+def parse_address(grpc_address: str) -> ParseResult:
+    if not grpc_address or not grpc_address.strip():
+        raise QwakGrpcAddressException(
+            "gRPC address must not be empty or whitespace.", grpc_address
+        )
+
+    trimmed_address: str = grpc_address.strip()
+    parsed_address: ParseResult = urlparse(
+        trimmed_address if "://" in trimmed_address else f"//{trimmed_address}"
+    )
+
+    return parsed_address
+
+
+def is_valid_hostname(hostname: str) -> bool:
+    """
+    Validate that the supplied hostname conforms to RFC-style label rules:
+    anchored pattern enforces full-string validation, negative lookahead/lookbehind block
+    leading or trailing hyphens per label, and each dot-separated label must be 1-63
+    alphanumeric/hyphen characters.
+
+    Args:
+        hostname (str): The hostname to validate.
+    Returns:
+        bool: True if the hostname is valid, False otherwise.
+    """
+    hostname_pattern: re.Pattern = re.compile(HOSTNAME_REGEX)
+    return bool(hostname_pattern.fullmatch(hostname))
+
+
 class SleepingPolicy(ABC):
     @abstractmethod
     def sleep(self, try_i: int):