qwak-sdk 0.5.102__py3-none-any.whl

qwak_sdk/commands/automations/automations_commands_group.py

@@ -0,0 +1,30 @@
+import logging
+
+import click
+
+from qwak_sdk.commands.automations.delete.ui import delete
+from qwak_sdk.commands.automations.executions.executions_commands_group import (
+    executions_command_group,
+)
+from qwak_sdk.commands.automations.list.ui import list_automations
+from qwak_sdk.commands.automations.register.ui import register
+
+logger = logging.getLogger(__name__)
+DELIMITER = "----------------------------------------"
+
+AUTOMATION = "automation"
+
+
+@click.group(
+    name="automations",
+    help="Commands for interacting with the Qwak Automations",
+)
+def automations_commands_group():
+    # Click commands group injection
+    pass
+
+
+automations_commands_group.add_command(delete)
+automations_commands_group.add_command(list_automations)
+automations_commands_group.add_command(register)
+automations_commands_group.add_command(executions_command_group)

qwak_sdk/commands/automations/delete/_logic.py

@@ -0,0 +1,6 @@
+from qwak.clients.automation_management.client import AutomationsManagementClient
+
+
+def delete_automation(automation_id: str):
+    client = AutomationsManagementClient()
+    return client.delete_automation(automation_id=automation_id)

qwak_sdk/commands/automations/delete/ui.py

@@ -0,0 +1,23 @@
+import click
+
+from qwak_sdk.commands.automations.delete._logic import delete_automation
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+
+
+@click.command(
+    "delete",
+    help="Delete an automation object",
+    cls=QwakCommand,
+)
+@click.option(
+    "--automation-id",
+    type=str,
+    metavar="ID",
+    help="The automation id to delete",
+)
+def delete(automation_id: str, **kwargs):
+    deleted = delete_automation(automation_id=automation_id)
+    if deleted:
+        print(f"Automation {automation_id} deleted successfully")
+    else:
+        print(f"Automation {automation_id} was not found")

qwak_sdk/commands/automations/executions/executions_commands_group.py

@@ -0,0 +1,14 @@
+import click
+
+from qwak_sdk.commands.automations.executions.list.ui import list_executions
+
+DELIMITER = "----------------------------------------"
+
+
+@click.group(name="executions", help=" automation executions")
+def executions_command_group():
+    # Click commands group injection
+    pass
+
+
+executions_command_group.add_command(list_executions)

qwak_sdk/commands/automations/executions/list/_logic.py

@@ -0,0 +1,8 @@
+from typing import List
+
+from qwak.automations.automation_executions import AutomationExecution
+from qwak.clients.automation_management.client import AutomationsManagementClient
+
+
+def execute_list_executions(automation_id: str) -> List[AutomationExecution]:
+    return AutomationsManagementClient().list_executions(automation_id)

qwak_sdk/commands/automations/executions/list/ui.py

@@ -0,0 +1,25 @@
+import click
+
+from qwak_sdk.commands.automations.executions.list._logic import execute_list_executions
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+
+DELIMITER = "----------------------------------------"
+
+
+@click.command(
+    "list",
+    help="List all executions for a specific automation",
+    cls=QwakCommand,
+)
+@click.option(
+    "--automation-id",
+    type=str,
+    metavar="ID",
+    required=True,
+    help="The automation id to list its executions",
+)
+def list_executions(automation_id: str, **kwargs):
+    executions_list = execute_list_executions(automation_id)
+    for execution in executions_list:
+        print(execution)
+        print(DELIMITER)

qwak_sdk/commands/automations/list/_logic.py

@@ -0,0 +1,36 @@
+from qwak.clients.automation_management.client import AutomationsManagementClient
+from tabulate import tabulate
+
+from qwak_sdk.commands._logic.tools import list_of_messages_to_json_str
+
+COLUMNS = ["Id", "Name", "Model id", "Type", "Status", "Last updated", "Updated by"]
+
+
+def execute_list_automations():
+    client = AutomationsManagementClient()
+    automations_list = client.list_automations()
+    data = []
+    for automation in automations_list:
+        data.append(
+            [
+                automation.id,
+                automation.name,
+                automation.model_id,
+                automation.action.__class__.__name__,
+                _map_automation_status(automation.is_enabled),
+                automation.create_audit.date,
+                automation.create_audit.user_id,
+            ]
+        )
+
+    return tabulate(data, headers=COLUMNS)
+
+
+def execute_list_json_automations() -> str:
+    client = AutomationsManagementClient()
+    automations = client.get_list_automations_from_server()
+    return list_of_messages_to_json_str(automations)
+
+
+def _map_automation_status(status: bool) -> str:
+    return "Active" if status else "Inactive"

qwak_sdk/commands/automations/list/ui.py

@@ -0,0 +1,21 @@
+import click
+
+from qwak_sdk.commands.automations.list._logic import (
+    execute_list_automations,
+    execute_list_json_automations,
+)
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+
+
+@click.command(
+    "list",
+    help="List all automations",
+    cls=QwakCommand,
+)
+@click.option("--json-format", is_flag=True, type=bool)
+def list_automations(json_format: bool, **kwargs):
+    if json_format:
+        click.echo(execute_list_json_automations())
+
+    else:
+        click.echo(execute_list_automations())

qwak_sdk/commands/automations/register/_logic.py

@@ -0,0 +1,43 @@
+from typing import List
+
+from qwak.automations import Automation
+from qwak.clients.automation_management.client import AutomationsManagementClient
+
+from qwak_sdk.inner.file_registry import extract_class_objects
+from qwak_sdk.inner.tools.cli_tools import ask_yesno
+from qwak_sdk.tools.utils import qwak_spinner
+
+DELIMITER = "----------------------------------------"
+
+
+def register_automations(qwak_python_files: List[str], force: bool):
+    """
+    Register Automation Entities Objects
+
+    Args:
+        qwak_python_files: a list of python files containing qwak package imports
+        force: to force
+    """
+    with qwak_spinner(
+        begin_text="Finding Automations to register", print_callback=print
+    ):
+        qwak_automations: List[Automation] = extract_class_objects(
+            qwak_python_files, Automation
+        )
+    client = AutomationsManagementClient()
+    print(f"Found {len(qwak_automations)} Automations")
+    for automation, source_file_path in qwak_automations:
+        existing_automation = client.get_automation_by_name(automation.name)
+        if existing_automation:
+            if ask_yesno(
+                f"Update existing Automation '{automation.name}' from source file '{source_file_path}'?",
+                force,
+            ):
+                client.update_automation(existing_automation.id, automation.to_proto())
+        else:
+            if ask_yesno(
+                f"Create new Automation '{automation.name}' from source file '{source_file_path}'?",
+                force,
+            ):
+                client.create_automation(automation.to_proto())
+    print(DELIMITER)

qwak_sdk/commands/automations/register/ui.py

@@ -0,0 +1,44 @@
+import os
+from pathlib import Path
+
+import click
+
+from qwak_sdk.commands.automations.register._logic import register_automations
+from qwak_sdk.inner.file_registry import list_qwak_python_files
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+from qwak_sdk.tools.utils import qwak_spinner
+
+
+@click.command(
+    "register",
+    help="Register all automations object under the given path. Registered "
+    "automations will be visible on the Qwak management platform after registration",
+    cls=QwakCommand,
+)
+@click.option(
+    "--path",
+    "-p",
+    type=click.Path(exists=True),
+    metavar="PATH",
+    help="Directory / module where qwak automations objects are stored",
+)
+@click.option(
+    "--force",
+    "-f",
+    default=False,
+    is_flag=True,
+    metavar="FLAG",
+    help="Force register all found qwak automations Store objects",
+)
+def register(path: Path, force: bool, **kwargs):
+    path = Path(path) if path else Path.cwd()
+    if path.is_file():
+        qwak_python_files = [(str(path), os.path.abspath(path))]
+    elif Path.is_dir(path):
+        with qwak_spinner(
+            begin_text="Recursively looking for python files in input dir",
+            print_callback=print,
+        ) as sp:
+            qwak_python_files = list_qwak_python_files(path, sp)
+
+    register_automations(qwak_python_files, force)

qwak_sdk/commands/feature_store/backfill/_logic.py

@@ -0,0 +1,140 @@
+from datetime import datetime, timezone
+from typing import Optional
+
+from _qwak_proto.qwak.feature_store.features.feature_set_pb2 import FeatureSet
+from _qwak_proto.qwak.feature_store.features.feature_set_service_pb2 import (
+    GetFeatureSetByNameResponse,
+)
+from croniter import croniter
+from qwak.clients.feature_store import FeatureRegistryClient
+from qwak.feature_store.execution.backfill import Backfill
+from qwak.feature_store.feature_sets.batch import BatchFeatureSet
+
+from qwak_sdk.inner.tools.cli_tools import ask_yesno
+from qwak_sdk.tools.colors import Color
+
+
+class _BackfillLogic:
+    featureset: FeatureSet
+    batch_featureset: BatchFeatureSet
+
+    def configure_window_backfill(
+        self,
+        featureset_name: str,
+        start_time: Optional[datetime] = None,
+        stop_time: Optional[datetime] = None,
+        comment: str = None,
+        cluster_template: str = None,
+    ) -> Optional[Backfill]:
+        """
+        Configures a window back-fill object. This type of backfill will not reset the entire featureset, but will be
+        performed from the requested start time to the requested stop time in ticks calculated according to the
+        feature sets scheduling policy.
+        @param featureset_name: feature set name to perform the backfill job for
+        @type featureset_name: str
+        @param start_time: backfill requested start date
+        @type start_time: optional[datetime]
+        @param stop_time:
+        @type stop_time: optional[datetime]
+        @param comment:
+        @type comment: str
+        @param cluster_template:
+        @type cluster_template: str
+        @return:
+        @rtype:
+        """
+        print(
+            f"{Color.BLUE} Backfilling from {start_time} to {stop_time}, "
+            f"the following ticks are going to be performed {Color.END}"
+        )
+        zipped_tick_strs = Backfill.generate_expected_ticks_repr(
+            scheduling_policy=self.batch_featureset.scheduling_policy,
+            start_time=start_time,
+            stop_time=stop_time,
+        )
+        print("\n".join(zipped_tick_strs))
+        if ask_yesno("", force=False):
+            backfill_execution = Backfill(
+                featureset_name=featureset_name,
+                comment=comment,
+                cluster_template=cluster_template,
+                start_time=start_time,
+                stop_time=stop_time,
+            )
+            return backfill_execution
+        return None
+
+    @staticmethod
+    def _get_featureset_by_name(featureset_name: str) -> GetFeatureSetByNameResponse:
+        registry_client = FeatureRegistryClient()
+
+        return registry_client.get_feature_set_by_name(feature_set_name=featureset_name)
+
+    def get_start_stop_times(
+        self,
+        requested_start_time: Optional[datetime] = None,
+        requested_stop_time: Optional[datetime] = None,
+    ) -> tuple:
+        start_time: datetime = (
+            requested_start_time or self.batch_featureset.backfill.start_date
+        ).replace(tzinfo=timezone.utc)
+        stop_time: datetime
+
+        if start_time and start_time > datetime.now().replace(tzinfo=timezone.utc):
+            stop_time = datetime.utcnow().replace(tzinfo=timezone.utc)
+        else:
+            stop_time = (
+                requested_stop_time or datetime.utcnow().replace(tzinfo=timezone.utc)
+            ).replace(tzinfo=timezone.utc)
+
+        return start_time, stop_time
+
+    def is_valid_input(
+        self,
+        reset_backfill: bool,
+        start_time: Optional[datetime],
+        stop_time: Optional[datetime],
+        featureset_name: str,
+    ) -> bool:
+        # Validate reset_backfill or start and stop times
+        if (not reset_backfill and not start_time and not stop_time) or (
+            reset_backfill and (start_time or stop_time)
+        ):
+            print(
+                f"{Color.RED} please specify either --reset-backfill or one or more of --start-time/--stop-time"
+            )
+            return False
+
+        # Validate featureset exists
+        get_featureset_resp = self._get_featureset_by_name(
+            featureset_name=featureset_name
+        )
+        if not get_featureset_resp or not get_featureset_resp.feature_set:
+            print(f"{Color.RED} Failed to retrieve featureset {featureset_name}")
+            return False
+
+        self.featureset = get_featureset_resp.feature_set
+        fs_type_attr: str = self.featureset.feature_set_definition.feature_set_spec.feature_set_type.WhichOneof(
+            "set_type"
+        )
+        # Validate featureset is batch v1
+        if fs_type_attr != "batch_feature_set_v1":
+            print(
+                f"{Color.RED} Feature set {featureset_name} is of type {fs_type_attr}. Only BatchV1 is supported."
+            )
+            return False
+
+        self.batch_featureset = BatchFeatureSet._from_proto(
+            self.featureset.feature_set_definition.feature_set_spec
+        )
+
+        # Validate scheduling policy
+        if not self.batch_featureset.scheduling_policy or not croniter.is_valid(
+            self.batch_featureset.scheduling_policy
+        ):
+            print(
+                f"{Color.RED} Feature set {featureset_name} has an invalid scheduling policy {self.batch_featureset.scheduling_policy}."
+            )
+            return False
+
+        return True

qwak_sdk/commands/feature_store/backfill/streaming/_logic.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+from typing import List, Tuple
+
+from qwak.feature_store.execution.streaming_backfill import StreamingAggregationBackfill
+from qwak.feature_store.feature_sets.streaming_backfill import StreamingBackfill
+
+from qwak_sdk.inner.file_registry import extract_class_objects
+from qwak_sdk.tools.utils import qwak_spinner
+
+DELIMITER = "----------------------------------------"
+
+
+def _trigger_streaming_backfills(backfill_definition_files: List[Tuple[str, str]]):
+    """
+    Trigger streaming aggregation backfills from Python files.
+
+    Args:
+        backfill_definition_files: List of tuples (module_name, file_path) containing qwak imports
+    """
+    with qwak_spinner(
+        begin_text="Finding streaming backfill definitions", print_callback=print
+    ):
+        streaming_backfills: List[
+            Tuple[StreamingBackfill, str]
+        ] = extract_class_objects(backfill_definition_files, StreamingBackfill)
+
+    print(f"👀 Found {len(streaming_backfills)} streaming backfill definition(s)")
+
+    print(DELIMITER)
+
+    # Trigger each backfill
+    for backfill_spec, source_file_path in streaming_backfills:
+        print(
+            f"Triggering backfill for featureset '{backfill_spec.featureset_name}' from '{source_file_path}'"
+        )
+
+        try:
+            # Create executor and trigger
+            executor = StreamingAggregationBackfill(
+                backfill_spec, Path(source_file_path)
+            )
+            execution_id = executor.trigger()
+
+            print(f"✅ Execution ID: {execution_id}")
+
+        except Exception as e:
+            print(f"❌ Failed: {str(e)}")
+            exit(1)
+
+        print(DELIMITER)

qwak_sdk/commands/feature_store/backfill/streaming/ui.py

@@ -0,0 +1,67 @@
+from pathlib import Path
+from typing import List, Tuple
+
+import click
+
+from qwak_sdk.commands.feature_store.backfill.streaming._logic import (
+    _trigger_streaming_backfills,
+)
+from qwak_sdk.inner.file_registry import list_qwak_python_files
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+from qwak_sdk.tools.utils import qwak_spinner
+
+
+@click.command(
+    "streaming",
+    cls=QwakCommand,
+    help="Trigger backfills for all streaming aggregation backfill definitions under the given path",
+)
+@click.option(
+    "--path",
+    "-p",
+    type=click.Path(exists=True, path_type=Path),
+    required=True,
+    help="Directory or file path containing backfill definitions (decorated with @streaming.backfill)",
+)
+def backfill_streaming(path: Path, **kwargs):
+    """
+    Trigger streaming aggregation backfills from a definition file.
+
+    The file should contain functions decorated with @streaming.backfill
+    that defines the backfill transformation and configuration.
+
+    Examples:
+        An example file can look like this:
+            from qwak.feature_store.feature_sets import streaming
+            from qwak.feature_store.feature_sets.streaming_backfill import BackfillDataSource
+            from datetime import datetime
+
+            @streaming.backfill(
+                feature_set_name="my_streaming_fs",
+                start_date=datetime(2023, 1, 1),
+                end_date=datetime(2023, 12, 1),
+                data_sources=[BackfillDataSource(data_source_name="batch_source")]
+            )
+
+            def my_backfill():
+                return SparkSqlTransformation("SELECT * FROM batch_source")
+
+        Usage:
+            qwak features backfill streaming --path backfill_definition.py
+    """
+    backfill_definition_files: List[Tuple[str, str]]
+
+    if Path.is_dir(path):
+        backfill_definition_files = _handle_directory(path)
+    else:
+        backfill_definition_files = [(str(path), str(path.absolute()))]
+
+    _trigger_streaming_backfills(backfill_definition_files)
+
+
+def _handle_directory(path: Path) -> list[tuple[str, str]]:
+    with qwak_spinner(
+        begin_text="Recursively looking for python files in input dir",
+        print_callback=print,
+    ) as sp:
+        return list_qwak_python_files(path, sp)

qwak_sdk/commands/feature_store/backfill/ui.py

@@ -0,0 +1,146 @@
+from datetime import datetime
+from typing import Optional
+
+import click
+from qwak.feature_store.execution.backfill import Backfill
+
+from qwak_sdk.commands.feature_store.backfill._logic import _BackfillLogic
+from qwak_sdk.commands.feature_store.backfill.streaming.ui import backfill_streaming
+from qwak_sdk.inner.tools.cli_tools import DefaultCommandGroup, QwakCommand, ask_yesno
+from qwak_sdk.tools.colors import Color
+
+
+@click.group(
+    cls=DefaultCommandGroup,
+    name="backfill",
+    default="batch",
+    help="Trigger a backfill process for Featuresets. If no command is provided, the default `batch` command is executed.",
+)
+def backfill():
+    """Backfill group with default batch_backfill command."""
+    pass
+
+
+# Default Command
+@click.command(
+    "batch", cls=QwakCommand, help="Trigger a backfill process for a Feature Set"
+)
+@click.option(
+    "--feature-set",
+    "--name",
+    help="Feature Set name to perform the backfill process for",
+    required=True,
+    type=click.STRING,
+)
+@click.option(
+    "--reset-backfill",
+    "--reset",
+    is_flag=True,
+    metavar="FLAG",
+    default=False,
+    help="Perform a complete reset of the featuresets data. "
+    "This will result in the deletion of the current existing data.",
+)
+@click.option(
+    "--start-time",
+    type=click.DateTime(),
+    default=None,
+    required=False,
+    help="The time from which the featuresets data should be backfilled in UTC. "
+    "Defaults to the featuresets configured backfill start time.",
+)
+@click.option(
+    "--stop-time",
+    type=click.DateTime(),
+    default=None,
+    required=False,
+    help="The time up until the featuresets data should be backfilled in UTC. Defaults to the current timestamp. "
+    "If the time provided is in the future, stop time will be rounded down to the current time. ",
+)
+@click.option(
+    "--cluster-template",
+    type=str,
+    default=None,
+    required=False,
+    help="Backfill resource configuration, expects a ClusterType size. "
+    "Defaults to the featureset resource configuration",
+)
+@click.option(
+    "--comment",
+    type=str,
+    default=None,
+    required=False,
+    help="Backfill job optional comment tag line",
+)
+def backfill_batch(
+    feature_set: str,
+    reset_backfill: bool,
+    start_time: Optional[datetime],
+    stop_time: Optional[datetime],
+    cluster_template: Optional[str],
+    comment: Optional[str],
+    **kwargs,
+):
+    backill_logic = _BackfillLogic()
+    if not backill_logic.is_valid_input(
+        reset_backfill=reset_backfill,
+        start_time=start_time,
+        stop_time=stop_time,
+        featureset_name=feature_set,
+    ):
+        return
+
+    if reset_backfill:
+        backfill_execution = _configure_backfill_triggered(
+            featureset_name=feature_set,
+            comment=comment,
+            cluster_template=cluster_template,
+        )
+    else:
+        start, stop = backill_logic.get_start_stop_times(
+            requested_start_time=start_time, requested_stop_time=stop_time
+        )
+        if start >= stop:
+            print(
+                f"{Color.RED} Stop time {stop.isoformat()} must be after start time {start.isoformat()}"
+            )
+            return
+
+        backfill_execution = backill_logic.configure_window_backfill(
+            featureset_name=feature_set,
+            start_time=start,
+            stop_time=stop,
+            comment=comment,
+            cluster_template=cluster_template,
+        )
+
+    if backfill_execution:
+        execution_id: str = backfill_execution.trigger_batch_backfill()
+        print(
+            f"{Color.BLUE}✅  Triggered backfill execution for featureset {feature_set}, "
+            f"execution id for follow up on is {execution_id}."
+        )
+        print(
+            f"To inquire the status of the execution run "
+            f"'qwak features execution-status --execution-id {execution_id}'"
+        )
+
+
+def _configure_backfill_triggered(
+    featureset_name: str, comment: str, cluster_template: str
+) -> Optional[Backfill]:
+    if ask_yesno(
+        f"You are about to remove and re-populate all data in {featureset_name}",
+        force=False,
+    ):
+        print(f"{Color.RED} - A backfill reset was triggered")
+        return Backfill(
+            featureset_name=featureset_name,
+            comment=comment,
+            cluster_template=cluster_template,
+        )
+    return None
+
+
+backfill.add_command(backfill_batch)
+backfill.add_command(backfill_streaming)