qwak-sdk 0.5.102__py3-none-any.whl

qwak_sdk/commands/feature_store/delete/_logic.py

@@ -0,0 +1,104 @@
+from _qwak_proto.qwak.feature_store.entities.entity_pb2 import Entity, EntityDefinition
+from _qwak_proto.qwak.feature_store.entities.entity_service_pb2 import (
+    GetEntityByNameResponse,
+)
+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 qwak.clients.feature_store import FeatureRegistryClient
+from qwak.exceptions import QwakException
+from qwak.feature_store._common.feature_set_utils import _INTERNAL_KEY_PREFIX
+
+from qwak_sdk.tools.colors import Color
+
+FEATURE_SET = "feature_set"
+ENTITY = "entity"
+DATA_SOURCE = "data_source"
+All = "all"
+
+
+def _handle_featureset_key_deletion(
+    entity_definition: EntityDefinition, registry_client: FeatureRegistryClient
+):
+    # If entity name start with the fs name
+    if entity_definition.entity_spec.name.startswith(_INTERNAL_KEY_PREFIX):
+        get_entity_resp: GetEntityByNameResponse = registry_client.get_entity_by_name(
+            entity_definition.entity_spec.name
+        )
+
+        # the entity exists
+        if get_entity_resp:
+            entity_id: str = entity_definition.entity_id
+            registered_entity: Entity = get_entity_resp.entity
+
+            # it has a no registered fs for it
+            if (
+                not registered_entity.feature_sets
+                or len(registered_entity.feature_sets) == 0
+            ):
+                try:
+                    # delete since it's a key and not an entity
+                    registry_client.delete_entity(entity_id=entity_id)
+                # we fail this silently
+                except QwakException:
+                    pass
+
+
+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_response: GetFeatureSetByNameResponse = (
+            registry_client.get_feature_set_by_name(name)
+        )
+        if feature_set_response:
+            feature_set: FeatureSet = feature_set_response.feature_set
+
+            # Delete fs regardless
+            registry_client.delete_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"
+            )
+
+            entity: EntityDefinition = (
+                feature_set.feature_set_definition.feature_set_spec.entity
+            )
+
+            _handle_featureset_key_deletion(
+                entity_definition=entity, registry_client=registry_client
+            )
+        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/execution/ui.py

@@ -0,0 +1,19 @@
+import click
+from qwak.feature_store.execution.execution_query import ExecutionQuery
+
+from qwak_sdk.inner.tools.cli_tools import QwakCommand
+
+
+@click.command(
+    "execution-status",
+    cls=QwakCommand,
+    help="Retrieve the current status of an execution (backfill/batch job ingestion)",
+)
+@click.option(
+    "--execution-id",
+    type=str,
+    required=True,
+    help="The id of the execution to retrieve the status for.",
+)
+def execution_status(execution_id: str, **kwargs):
+    print(ExecutionQuery.get_execution_status_message(execution_id=execution_id))

qwak_sdk/commands/feature_store/feature_store_command_group.py

@@ -0,0 +1,29 @@
+import click
+
+from qwak_sdk.commands.feature_store.backfill.ui import backfill
+from qwak_sdk.commands.feature_store.delete.ui import delete_fs_object
+from qwak_sdk.commands.feature_store.execution.ui import execution_status
+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)
+feature_store_commands_group.add_command(backfill)
+feature_store_commands_group.add_command(execution_status)

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}")