Sammyt/add snapshot date param (#310)

* refactor: indent style the import, because the editor says so today

* feature: add input_date param in snapshot check command

* refactor: use same logic for today

tools/driftdb/driftdb/cli/common.py

@@ -1,7 +1,9 @@
+from datetime import datetime
+
 import inquirer
 import pandas as pd
 import typer
-from typing_extensions import List
+from typing_extensions import List, Optional
 
 
 def prompt_from_list(prompt: str, choices: list):
@@ -24,8 +26,22 @@ def dbt_adapter_query(
     data = {column_name: table.columns[column_name].values() for column_name in table.column_names}
     return pd.DataFrame(data)
 
+def find_date_starting_with(dates, input_date):
+    for date in dates:
+        if date.startswith(input_date):
+            return date
+    return None
+
+def get_user_date_selection(dates: List[str], input_date: str) -> Optional[str]:
+    if input_date is not None:
+        if input_date == "today":
+            input_date = datetime.today().date().strftime("%Y-%m-%d")
+
+
+        print("dates",input_date)
+        matching_date = find_date_starting_with(dates, input_date)
+        return matching_date
 
-def get_user_date_selection(dates: List[str]) -> str:
     page_size = 10
     page = 0
 

tools/driftdb/driftdb/cli/snapshot.py

@@ -1,5 +1,4 @@
 import base64
-import json
 import os
 import webbrowser
 
@@ -8,7 +7,8 @@
 import typer
 
 from ..alerting.handlers import alert_drift_handler
-from ..dbt.snapshot import get_snapshot_dates, get_snapshot_diff, get_snapshot_nodes
+from ..dbt.snapshot import (get_snapshot_dates, get_snapshot_diff,
+                            get_snapshot_nodes)
 from ..dbt.snapshot_to_drift import convert_snapshot_to_drift_summary
 from ..logger import get_logger
 from ..user_defined_function import import_user_defined_function
@@ -51,10 +51,14 @@ def show(snapshot_id: str = typer.Option(None, help="id of your snapshot")):
 
 
 @app.command()
-def check(snapshot_id: str = typer.Option(None, help="id of your snapshot")):
+def check(snapshot_id: str = typer.Option(None, help="id of your snapshot"), date: str = typer.Option(None, help="date of your snapshot")):
     snapshot_node = get_or_prompt_snapshot_node(snapshot_id, get_snapshot_nodes())
     handler = get_snapshot_handler(snapshot_node)
-    snapshot_date = get_user_date_selection(get_snapshot_dates(snapshot_node))
+    snapshot_date = get_user_date_selection(get_snapshot_dates(snapshot_node), date)
+
+    if snapshot_date is None:
+        typer.echo("No snapshot data for selected date. Exiting.")
+        raise typer.Exit(code=1)
 
     print(f"Getting {snapshot_node['unique_id']} for {snapshot_date}.")
 

tools/driftdb/driftdb/version.py

@@ -1 +1 @@
-version = "0.1.3-a.9"
+version = "0.1.3-a.10"