qwak-sdk 0.1.0__py3-none-any.whl → 0.2.21rc0__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,8 @@
+from qwak.clients.automation_management.client import AutomationsManagementClient
+
+DELIMITER = "----------------------------------------"
+
+
+def execute_list_automations():
+    client = AutomationsManagementClient()
+    return client.list_automations()

qwak_sdk/commands/automations/list/ui.py

@@ -0,0 +1,21 @@
+from typing import List
+
+import click
+from qwak.automations import Automation
+
+from qwak_sdk.commands.automations.list._logic import execute_list_automations
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+
+DELIMITER = "----------------------------------------"
+
+
+@click.command(
+    "list",
+    help="List all automations",
+    cls=QwakCommand,
+)
+def list_automations(**kwargs):
+    automations_list: List[Automation] = execute_list_automations()
+    for automation in automations_list:
+        print(automation)
+        print(DELIMITER)

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/delete/_logic.py

@@ -0,0 +1,52 @@
+from qwak_sdk.tools.colors import Color
+
+FEATURE_SET = "feature_set"
+ENTITY = "entity"
+DATA_SOURCE = "data_source"
+All = "all"
+
+
+def _inner_delete(name, object, registry_client):
+    if object == ENTITY:
+        feature_store_entity = registry_client.get_entity_by_name(name)
+        _delete_object(
+            lambda e: registry_client.delete_entity(
+                e.entity.entity_definition.entity_id
+            ),
+            feature_store_entity,
+            "entity",
+            name,
+        )
+    if object == FEATURE_SET:
+        feature_set = registry_client.get_feature_set_by_name(name)
+        if feature_set:
+            registry_client.delete_feature_set(
+                feature_set.feature_set.feature_set_definition.feature_set_id
+            )
+            print(
+                f"{Color.GREEN}Feature Set '{name}' deletion request is being handled by the service"
+            )
+        else:
+            print(
+                f"{Color.RED}Could not find specified feature set named '{name}' to delete{Color.END}"
+            )
+    if object == DATA_SOURCE:
+        feature_store_ds = registry_client.get_data_source_by_name(name)
+        _delete_object(
+            lambda e: registry_client.delete_data_source(
+                e.data_source.data_source_definition.data_source_id
+            ),
+            feature_store_ds,
+            "data source",
+            name,
+        )
+
+
+def _delete_object(delete_func, _object, _type, name):
+    if _object:
+        delete_func(_object)
+        print(f"{Color.GREEN}Successfully deleted the {_type} named '{name}'")
+    else:
+        print(
+            f"{Color.RED}Could not find specified {_type} named '{name}' to delete{Color.END}"
+        )

qwak_sdk/commands/feature_store/delete/ui.py

@@ -0,0 +1,40 @@
+import click
+from qwak.clients.feature_store import FeatureRegistryClient
+
+from qwak_sdk.commands.feature_store.delete._logic import (
+    DATA_SOURCE,
+    ENTITY,
+    FEATURE_SET,
+    _inner_delete,
+)
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+
+
+@click.command(
+    "delete",
+    cls=QwakCommand,
+    help="Delete by name a feature store object - a feature set, entity, or data source",
+)
+@click.option(
+    "--feature-set",
+    "object",
+    flag_value=FEATURE_SET,
+    default=True,
+    help="Delete a feature set [default option]",
+)
+@click.option(
+    "--data-source", "object", flag_value=DATA_SOURCE, help="Delete a data source"
+)
+@click.option("--entity", "object", flag_value=ENTITY, help="Delete an entity")
+@click.argument("name")
+def delete_fs_object(object, name, **kwargs):
+    """
+    Delete Feature Store object
+
+    Args:
+        object: type of object to delete
+        name: name of the object to delete
+    """
+
+    registry_client = FeatureRegistryClient()
+    _inner_delete(name, object, registry_client)

qwak_sdk/commands/feature_store/feature_store_command_group.py

@@ -0,0 +1,25 @@
+import click
+
+from qwak_sdk.commands.feature_store.delete.ui import delete_fs_object
+from qwak_sdk.commands.feature_store.list.ui import list_feature_sets
+from qwak_sdk.commands.feature_store.pause.ui import pause_feature_set
+from qwak_sdk.commands.feature_store.register.ui import register_fs_objects
+from qwak_sdk.commands.feature_store.resume.ui import resume_feature_set
+from qwak_sdk.commands.feature_store.trigger.ui import trigger_feature_set
+
+
+@click.group(
+    name="features",
+    help="Commands for interacting with the Qwak Feature Store",
+)
+def feature_store_commands_group():
+    # Click commands group injection
+    pass
+
+
+feature_store_commands_group.add_command(delete_fs_object)
+feature_store_commands_group.add_command(list_feature_sets)
+feature_store_commands_group.add_command(pause_feature_set)
+feature_store_commands_group.add_command(resume_feature_set)
+feature_store_commands_group.add_command(trigger_feature_set)
+feature_store_commands_group.add_command(register_fs_objects)

qwak_sdk/commands/feature_store/list/ui.py

@@ -0,0 +1,140 @@
+from datetime import datetime
+from typing import List, Tuple
+
+import click
+from _qwak_proto.qwak.feature_store.features.feature_set_pb2 import (
+    FeaturesetSchedulingState,
+    FeatureStatus,
+)
+from _qwak_proto.qwak.feature_store.features.feature_set_service_pb2 import (
+    GetFeaturesetSchedulingStateResponse,
+    ListFeatureSetsResponse,
+)
+from qwak.clients.feature_store import FeatureRegistryClient
+from tabulate import tabulate
+
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+
+
+def _pts_to_datetime(pts) -> datetime:
+    return datetime.fromtimestamp(pts.seconds + pts.nanos / 1e9)
+
+
+def _datetime_to_str(dt: datetime) -> str:
+    return dt.strftime("%Y-%b-%d %H:%M:%S")
+
+
+def get_featureset_scheduling_state_as_str(
+    registry_client: FeatureRegistryClient,
+    featureset_name: str,
+) -> str:
+    try:
+        resp: GetFeaturesetSchedulingStateResponse = (
+            registry_client.get_feature_set_scheduling_state(
+                featureset_name=featureset_name
+            )
+        )
+        if resp.state == FeaturesetSchedulingState.SCHEDULING_STATE_PAUSED:
+            return "Paused"
+        elif resp.state == FeaturesetSchedulingState.SCHEDULING_STATE_ENABLED:
+            return "Enabled"
+        else:
+            return "Unknown"
+    except Exception:
+        return "Unknown"
+
+
+def _list_feature_sets_verbose(
+    resp: ListFeatureSetsResponse, registry_client: FeatureRegistryClient
+) -> Tuple[List[str], List[str]]:
+    columns = [
+        "Name",
+        "Type",
+        "Status",
+        "Scheduling state",
+        "Owner",
+        "Description",
+        "Creation date",
+        "Last updated",
+        "Featureset id",
+    ]
+    data = []
+    for feature_set in resp.feature_families:
+        definition = feature_set.feature_set_definition
+        spec = definition.feature_set_spec
+        metadata = feature_set.metadata
+        if not (spec.feature_set_type.WhichOneof("set_type")):
+            feature_set_type = ""
+        else:
+            feature_set_type = (
+                spec.feature_set_type.WhichOneof("set_type")
+                .replace("_", " ")
+                .capitalize()
+            )
+
+        sched_state: str = "Enabled"
+        if spec.feature_set_type.WhichOneof("set_type") in {
+            "batch_feature_set",
+            "batch_feature_set_v1",
+        }:
+            sched_state = get_featureset_scheduling_state_as_str(
+                registry_client, spec.name
+            )
+
+        data.append(
+            [
+                spec.name,
+                feature_set_type,
+                FeatureStatus.Name(definition.status),
+                sched_state,
+                spec.metadata.owner,
+                spec.metadata.description,
+                _datetime_to_str(_pts_to_datetime(metadata.created_at)),
+                _datetime_to_str(_pts_to_datetime(metadata.last_modified_at)),
+                definition.feature_set_id,
+            ]
+        )
+
+    return (data, columns)
+
+
+def _list_feature_sets(
+    resp: ListFeatureSetsResponse, registry_client: FeatureRegistryClient
+) -> Tuple[List[str], List[str]]:
+    desired_columns = [
+        "Name",
+        "Type",
+        "Status",
+        "Scheduling state",
+        "Last updated",
+    ]
+    verbose_data, verbose_columns = _list_feature_sets_verbose(resp, registry_client)
+    col_indexes = [verbose_columns.index(c) for c in desired_columns]
+
+    def project_lst(lst: List[str]) -> List[str]:
+        return [lst[idx] for idx in col_indexes]
+
+    data = [project_lst(d) for d in verbose_data]
+
+    return (data, desired_columns)
+
+
+@click.command("list", cls=QwakCommand, help="List registered feature sets")
+@click.option(
+    "--verbose",
+    "-v",
+    default=False,
+    is_flag=True,
+    metavar="FLAG",
+    help="Verbose output",
+)
+def list_feature_sets(verbose: bool, **kwargs):
+    registry_client = FeatureRegistryClient()
+    response = registry_client.list_feature_sets()
+
+    if verbose:
+        data, columns = _list_feature_sets_verbose(response, registry_client)
+    else:
+        data, columns = _list_feature_sets(response, registry_client)
+
+    print(tabulate(data, headers=columns))

qwak_sdk/commands/feature_store/pause/ui.py

@@ -0,0 +1,18 @@
+import click
+from qwak.clients.feature_store import FeatureRegistryClient
+from qwak.exceptions import QwakException
+
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+from qwak_sdk.tools.colors import Color
+
+
+@click.command("pause", cls=QwakCommand, help="Pause a running feature set")
+@click.argument("name")
+def pause_feature_set(name, **kwargs):
+    try:
+        FeatureRegistryClient().pause_feature_set(feature_set_name=name)
+    except Exception as e:
+        print(f"{Color.RED} Failed to pause feature set {name} {Color.END}")
+        raise QwakException(f"Failed to pause feature set {name}") from e
+
+    print(f"{Color.GREEN}Successfully paused feature set {Color.YELLOW}{name}")