feat: select dates for workflow.filter

Signed-off-by: Yves Bastide <yves@botify.com>

Author: ybastide

simpleflow/command.py

@@ -6,6 +6,7 @@
 import sys
 import time
 from contextlib import contextmanager
+from datetime import datetime
 from typing import TYPE_CHECKING
 from uuid import uuid4
 
@@ -32,6 +33,14 @@
 
     from simpleflow.swf.mapper.models.workflow import WorkflowType
 
+TIMESTAMP_FORMATS = [
+    "%Y-%m-%d",
+    "%Y-%m-%dT%H:%M:%S%z",
+    "%Y-%m-%d %H:%M:%S%z",
+    "%Y-%m-%dT%H:%M:%S",
+    "%Y-%m-%d %H:%M:%S",
+]
+
 
 def comma_separated_list(value):
     """
@@ -41,7 +50,7 @@ def comma_separated_list(value):
 
 
 @click.group()
-@click.option("--format")
+@click.option("--format", envvar="SIMPLEFLOW_FORMAT")
 @click.option("--header/--no-header", default=False)
 @click.option(
     "--color",
@@ -51,6 +60,8 @@ def comma_separated_list(value):
 @click.version_option(version=__version__)
 @click.pass_context
 def cli(ctx, header, format, color):
+    if format == "prettyjson":
+        format, header = "json", True
     ctx.params["format"] = format
     ctx.params["header"] = header
     log.color_mode = color
@@ -345,28 +356,60 @@ def list_workflows(ctx, domain, status, started_since):
     type=click.Choice(["open", "closed"]),
     help="Open/Closed",
 )
-@click.option("--tag", default=None, help="Tag (multiple option).")  # , multiple=True
+@click.option("--tag", default=None, help="Tags (comma-separated).")  # , multiple=True
 @click.option("--workflow-id", default=None, help="Workflow ID.")
 @click.option("--workflow-type-name", default=None, help="Workflow Name.")
 @click.option("--workflow-type-version", default=None, help="Workflow Version (name needed).")
 @click.option("--started-since", "-d", default=30, show_default=True, help="Started since N days.")
+@click.option("--from-date", default=None, type=click.DateTime(formats=TIMESTAMP_FORMATS), help="From datetime.")
+@click.option("--to-date", default=None, type=click.DateTime(formats=TIMESTAMP_FORMATS), help="To datetime.")
 @click.pass_context
 def filter_workflows(
     ctx,
-    domain,
-    status,
-    tag,
-    workflow_id,
-    workflow_type_name,
-    workflow_type_version,
-    started_since,
+    domain: str,
+    status: str,
+    tag: str | None,
+    workflow_id: str | None,
+    workflow_type_name: str | None,
+    workflow_type_version: str | None,
+    started_since: int | None,
+    from_date: datetime | None,
+    to_date: datetime | None,
 ):
     status = status.upper()
     kwargs = {}
     if status == simpleflow.swf.mapper.models.workflow.WorkflowExecution.STATUS_OPEN:
-        kwargs["oldest_date"] = started_since
+        if from_date:
+            kwargs["oldest_date"] = from_date
+            kwargs["latest_date"] = to_date
+        else:
+            kwargs["oldest_date"] = started_since
+    else:
+        if from_date:
+            kwargs["start_oldest_date"] = from_date
+            kwargs["start_latest_date"] = to_date
+        else:
+            kwargs["start_oldest_date"] = started_since
+
+    if os.isatty(sys.stderr.fileno()):
+        spin_marks = ["⠏", "⠛", "⠹", "⠼", "⠶", "⠧"]  # from Google's googlecloudsdk.core
+        success = f"{log.GREEN}✓{log.END}" if log.color_mode != log.ColorModes.NEVER else "✓"
+
+        class Counter:
+            def __init__(self):
+                self.total = 0
+
+        counter = Counter()
+
+        def cb(*_args, loop_number: int, response: dict[str, Any] | None, **_kwargs):
+            if response:
+                executions = len(response.get("executionInfos", []))
+                print(f"\r{spin_marks[loop_number % len(spin_marks)]} {counter.total}", file=sys.stderr, end="")
+                counter.total += executions
+            else:
+                print(f"\r{success} {counter.total}", file=sys.stderr)
     else:
-        kwargs["start_oldest_date"] = started_since
+        cb = None
     print(
         with_format(ctx)(helpers.filter_workflow_executions)(
             domain,
@@ -375,6 +418,7 @@ def filter_workflows(
             workflow_id=workflow_id,
             workflow_type_name=workflow_type_name,
             workflow_type_version=workflow_type_version,
+            callback=cb,
             **kwargs,
         )
     )

simpleflow/swf/helpers.py

@@ -98,6 +98,7 @@ def filter_workflow_executions(
     workflow_type_name,
     workflow_type_version,
     *args,
+    callback=None,
     **kwargs,
 ):
     domain = simpleflow.swf.mapper.models.Domain(domain_name)
@@ -109,6 +110,7 @@ def filter_workflow_executions(
         workflow_type_name,
         workflow_type_version,
         *args,
+        callback=callback,
         **kwargs,
     )
 

simpleflow/swf/mapper/core.py

@@ -24,6 +24,16 @@
 DEFAULT_AWS_REGION = "us-east-1"
 
 
+def convert_timestamp(dt: int | datetime | None) -> datetime | None:
+    if dt is None:
+        return None
+    if isinstance(dt, datetime):
+        return dt
+    if isinstance(dt, int):
+        return datetime.fromtimestamp(dt)
+    raise TypeError(f"Invalid timestamp type: {type(dt).__name__}")
+
+
 class ConnectedSWFObject:
     """Authenticated object interface"""
 
@@ -56,8 +66,8 @@ def __init__(self, *args, **kwargs):
     def list_open_workflow_executions(
         self,
         domain: str,
-        oldest_date: int,  # timestamp
-        latest_date: int | None = None,  # timestamp
+        oldest_date: int | datetime,  # timestamp
+        latest_date: int | datetime | None = None,  # timestamp
         tag: str | None = None,
         workflow_id: str | None = None,
         workflow_name: str | None = None,
@@ -69,8 +79,8 @@ def list_open_workflow_executions(
         kwargs = {
             "domain": domain,
             "startTimeFilter": {
-                "oldestDate": datetime.fromtimestamp(oldest_date),
-                "latestDate": datetime.fromtimestamp(latest_date) if latest_date is not None else None,
+                "oldestDate": convert_timestamp(oldest_date),
+                "latestDate": convert_timestamp(latest_date),
             },
             "nextPageToken": next_page_token,
             "maximumPageSize": maximum_page_size,
@@ -98,10 +108,10 @@ def list_open_workflow_executions(
     def list_closed_workflow_executions(
         self,
         domain: str,
-        start_latest_date: int | None = None,  # timestamp
-        start_oldest_date: int | None = None,  # timestamp
-        close_latest_date: int | None = None,  # timestamp
-        close_oldest_date: int | None = None,  # timestamp
+        start_latest_date: int | datetime | None = None,  # timestamp
+        start_oldest_date: int | datetime | None = None,  # timestamp
+        close_latest_date: int | datetime | None = None,  # timestamp
+        close_oldest_date: int | datetime | None = None,  # timestamp
         close_status: str | None = None,
         tag: str | None = None,
         workflow_id: str | None = None,
@@ -119,13 +129,13 @@ def list_closed_workflow_executions(
         }
         if start_oldest_date is not None:
             kwargs["startTimeFilter"] = {
-                "oldestDate": datetime.fromtimestamp(start_oldest_date),
-                "latestDate": datetime.fromtimestamp(start_latest_date) if start_latest_date is not None else None,
+                "oldestDate": convert_timestamp(start_oldest_date),
+                "latestDate": convert_timestamp(start_latest_date),
             }
         if close_oldest_date is not None:
             kwargs["closeTimeFilter"] = {
-                "oldestDate": datetime.fromtimestamp(close_oldest_date),
-                "latestDate": datetime.fromtimestamp(close_latest_date) if close_latest_date is not None else None,
+                "oldestDate": convert_timestamp(close_oldest_date),
+                "latestDate": convert_timestamp(close_latest_date),
             }
         if close_status:
             kwargs["closeStatusFilter"] = {

simpleflow/swf/mapper/querysets/workflow.py

@@ -55,12 +55,21 @@ def domain(self, value: Domain):
     def _list(self, *args, **kwargs):
         raise NotImplementedError
 
-    def _list_items(self, *args, **kwargs):
+    def _list_items(self, *args, callback=None, **kwargs):
         response = {"nextPageToken": None}
+        loop_number = 0
+        if callback:
+            callback(loop_number=loop_number, response=response)
         while "nextPageToken" in response:
             response = self._list(*args, next_page_token=response["nextPageToken"], **kwargs)
 
             yield from response[self._infos_plural]
+            loop_number += 1
+            if callback:
+                callback(loop_number=loop_number, response=response)
+
+        if callback:
+            callback(loop_number=loop_number, response=None)
 
 
 class WorkflowTypeQuerySet(BaseWorkflowQuerySet):
@@ -405,10 +414,10 @@ def _is_valid_status_param(self, status, param):
         statuses = {
             WorkflowExecution.STATUS_OPEN: {"oldest_date", "latest_date"},
             WorkflowExecution.STATUS_CLOSED: {
-                "start_latest_date",
                 "start_oldest_date",
-                "close_latest_date",
+                "start_latest_date",
                 "close_oldest_date",
+                "close_latest_date",
                 "close_status",
             },
         }
@@ -501,6 +510,7 @@ def filter(
         workflow_id=None,
         workflow_type_name=None,
         workflow_type_version=None,
+        callback=None,
         *args,
         **kwargs,
     ):
@@ -591,10 +601,13 @@ def filter(
 
         # Compute a timestamp from the delta in days we got from params
         # If oldest_date is blank at this point, it's because we didn't want
-        # it, so let's leave it blank and assume the user provided an other
+        # it, so let's leave it blank and assume the user provided another
         # time filter.
         if oldest_date:
-            start_oldest_date = int(datetime_timestamp(past_day(oldest_date)))
+            if isinstance(oldest_date, int):
+                start_oldest_date = int(datetime_timestamp(past_day(oldest_date)))
+            else:
+                start_oldest_date = oldest_date
         else:
             start_oldest_date = None
 
@@ -609,6 +622,7 @@ def filter(
                 workflow_version=workflow_type_version,
                 start_oldest_date=start_oldest_date,
                 tag=tag,
+                callback=callback,
                 **kwargs,
             )
         ]