frogml-cli 0.0.1__py3-none-any.whl

frogml_cli/commands/automations/register/ui.py

@@ -0,0 +1,44 @@
+import os
+from pathlib import Path
+
+import click
+
+from frogml_cli.commands.automations.register._logic import register_automations
+from frogml_cli.inner.file_registry import list_qwak_python_files
+from frogml_cli.inner.tools.cli_tools import FrogMLCommand
+from frogml_cli.tools.utils import frogml_spinner
+
+
+@click.command(
+    "register",
+    help="Register all automations object under the given path. Registered "
+    "automations will be visible on the Frogml management platform after registration",
+    cls=FrogMLCommand,
+)
+@click.option(
+    "--path",
+    "-p",
+    type=click.Path(exists=True),
+    metavar="PATH",
+    help="Directory / module where frogml automations objects are stored",
+)
+@click.option(
+    "--force",
+    "-f",
+    default=False,
+    is_flag=True,
+    metavar="FLAG",
+    help="Force register all found frogml 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 frogml_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)

frogml_cli/commands/config/add/ui.py

@@ -0,0 +1,62 @@
+from typing import Optional
+
+import click
+from frogml.core.inner.di_configuration import UserAccountConfiguration
+from frogml.core.inner.di_configuration.account import UserAccount
+
+from frogml_cli.inner.tools.cli_tools import FrogMLCommand
+
+
+@click.command("add", cls=FrogMLCommand, help="Adds a server configuration.")
+@click.option("--url", metavar="BASE_URL", required=False, help="Artifactory base url")
+@click.option(
+    "--username", metavar="USERNAME", required=False, help="The user's username"
+)
+@click.option(
+    "--password", metavar="PASSWORD", required=False, help="The user's password"
+)
+@click.option(
+    "--access-token",
+    metavar="ACCESS_TOKEN",
+    required=False,
+    help="Access token to authenticate",
+)
+@click.option(
+    "--interactive",
+    metavar="INTERACTIVE",
+    required=False,
+    help="Login with interactive flow",
+    default=False,
+)
+@click.option(
+    "--server-id",
+    metavar="SERVER_ID",
+    required=False,
+    help="The server id, if not provided, the server id will be auto generated as uuid.",
+)
+def add_config(
+    url: Optional[str],
+    username: Optional[str],
+    password: Optional[str],
+    access_token: Optional[str],
+    server_id: Optional[str],
+    interactive: Optional[bool],
+    **_,
+):
+    if url is None and not interactive:
+        raise click.BadOptionUsage(
+            option_name="config add",
+            message="One of --interactive or --url must be provided",
+        )
+    account_config = UserAccountConfiguration()
+    account_config.configure_user(
+        UserAccount(
+            url=url,
+            username=username,
+            password=password,
+            token=access_token,
+            anonymous=False,
+            is_interactive=interactive,
+            server_id=server_id,
+        )
+    )

frogml_cli/commands/config/config_commands_group.py

@@ -0,0 +1,11 @@
+import click
+
+from frogml_cli.commands.config.add.ui import add_config
+
+
+@click.group(name="config", help="Configure the FrogML environment")
+def config_commands_group():
+    pass
+
+
+config_commands_group.add_command(add_config)

frogml_cli/commands/feature_store/backfill/_logic.py

@@ -0,0 +1,140 @@
+from datetime import datetime, timezone
+from typing import Optional
+
+from croniter import croniter
+from frogml._proto.qwak.feature_store.features.feature_set_pb2 import FeatureSet
+from frogml._proto.qwak.feature_store.features.feature_set_service_pb2 import (
+    GetFeatureSetByNameResponse,
+)
+from frogml.core.clients.feature_store import FeatureRegistryClient
+from frogml.core.feature_store.execution.backfill import Backfill
+from frogml.feature_store.feature_sets.batch import BatchFeatureSet
+
+from frogml_cli.inner.tools.cli_tools import ask_yesno
+from frogml_cli.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

frogml_cli/commands/feature_store/backfill/ui.py

@@ -0,0 +1,129 @@
+from datetime import datetime
+from typing import Optional
+
+import click
+from frogml.core.feature_store.execution.backfill import Backfill
+
+from frogml_cli.commands.feature_store.backfill._logic import _BackfillLogic
+from frogml_cli.inner.tools.cli_tools import FrogMLCommand, ask_yesno
+from frogml_cli.tools.colors import Color
+
+
+@click.command(
+    "backfill", cls=FrogMLCommand, 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(
+    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"'frogml 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

frogml_cli/commands/feature_store/delete/_logic.py

@@ -0,0 +1,107 @@
+from frogml._proto.qwak.feature_store.entities.entity_pb2 import (
+    Entity,
+    EntityDefinition,
+)
+from frogml._proto.qwak.feature_store.entities.entity_service_pb2 import (
+    GetEntityByNameResponse,
+)
+from frogml._proto.qwak.feature_store.features.feature_set_pb2 import FeatureSet
+from frogml._proto.qwak.feature_store.features.feature_set_service_pb2 import (
+    GetFeatureSetByNameResponse,
+)
+from frogml.core.clients.feature_store import FeatureRegistryClient
+from frogml.core.exceptions import FrogmlException
+from frogml.feature_store._common.feature_set_utils import _INTERNAL_KEY_PREFIX
+
+from frogml_cli.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 FrogmlException:
+                    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}"
+        )

frogml_cli/commands/feature_store/delete/ui.py

@@ -0,0 +1,40 @@
+import click
+from frogml.core.clients.feature_store import FeatureRegistryClient
+
+from frogml_cli.commands.feature_store.delete._logic import (
+    DATA_SOURCE,
+    ENTITY,
+    FEATURE_SET,
+    _inner_delete,
+)
+from frogml_cli.inner.tools.cli_tools import FrogMLCommand
+
+
+@click.command(
+    "delete",
+    cls=FrogMLCommand,
+    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)

frogml_cli/commands/feature_store/execution/ui.py

@@ -0,0 +1,19 @@
+import click
+from frogml.core.feature_store.execution.execution_query import ExecutionQuery
+
+from frogml_cli.inner.tools.cli_tools import FrogMLCommand
+
+
+@click.command(
+    "execution-status",
+    cls=FrogMLCommand,
+    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))

frogml_cli/commands/feature_store/feature_store_command_group.py

@@ -0,0 +1,29 @@
+import click
+
+from frogml_cli.commands.feature_store.backfill.ui import backfill
+from frogml_cli.commands.feature_store.delete.ui import delete_fs_object
+from frogml_cli.commands.feature_store.execution.ui import execution_status
+from frogml_cli.commands.feature_store.list.ui import list_feature_sets
+from frogml_cli.commands.feature_store.pause.ui import pause_feature_set
+from frogml_cli.commands.feature_store.register.ui import register_fs_objects
+from frogml_cli.commands.feature_store.resume.ui import resume_feature_set
+from frogml_cli.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)