Merge pull request #125 from andrewm4894/add-delete-jobs

Add delete job functionality with SQL template and scheduling

Author: andrewm4894

anomstack/external/duckdb/duckdb.py

@@ -63,3 +63,31 @@ def save_df_duckdb(df: pd.DataFrame, table_key: str) -> pd.DataFrame:
         query(connection=conn, query=f"CREATE TABLE {table_key} AS SELECT * FROM df")
 
     return df
+
+
+def run_sql_duckdb(sql: str) -> None:
+    """
+    Execute a non-returning SQL statement in DuckDB.
+
+    Args:
+        sql (str): The SQL statement to execute.
+
+    Returns:
+        None
+    """
+    logger = get_dagster_logger()
+
+    duckdb_path = os.environ.get("ANOMSTACK_DUCKDB_PATH", "tmpdata/anomstack.db")
+    logger.info(f"duckdb_path: {duckdb_path}")
+
+    os.makedirs(os.path.dirname(duckdb_path), exist_ok=True)
+
+    conn = connect(duckdb_path)
+
+    try:
+        query(connection=conn, query=sql)
+    except Exception as e:
+        logger.error(f"Error executing SQL statement in DuckDB: {e}")
+        raise
+    finally:
+        conn.close()

anomstack/external/sqlite/sqlite.py

@@ -90,3 +90,47 @@ def save_df_sqlite(df: pd.DataFrame, table_key: str) -> pd.DataFrame:
                 raise
     # If all retries fail, raise an error
     raise sqlite3.OperationalError("Database is locked after multiple attempts.")
+
+
+def run_sql_sqlite(sql: str) -> None:
+    """
+    Execute a non-returning SQL statement in SQLite with retry logic.
+
+    Args:
+        sql (str): The SQL statement to execute.
+
+    Returns:
+        None
+    """
+    logger = get_dagster_logger()
+    sqlite_path = os.environ.get("ANOMSTACK_SQLITE_PATH", "tmpdata/anomstack.db")
+    logger.info(f"sqlite_path: {sqlite_path}")
+    os.makedirs(os.path.dirname(sqlite_path), exist_ok=True)
+
+    attempt = 0
+    while attempt < MAX_RETRIES:
+        try:
+            conn = sqlite3.connect(sqlite_path)
+            cursor = conn.cursor()
+            cursor.execute(sql)
+            conn.commit()
+            cursor.close()
+            conn.close()
+            return
+        except sqlite3.OperationalError as e:
+            if "database is locked" in str(e):
+                attempt += 1
+                logger.warning(
+                    f"Database is locked; attempt {attempt} of {MAX_RETRIES}. "
+                    f"Retrying in {RETRY_DELAY} seconds..."
+                )
+                time.sleep(RETRY_DELAY)
+            else:
+                logger.error(f"Error executing SQL statement: {e}")
+                raise
+        finally:
+            if 'conn' in locals():
+                conn.close()
+
+    # If all retries fail, raise an error
+    raise sqlite3.OperationalError("Database is locked after multiple attempts.")

anomstack/jobs/delete.py

@@ -0,0 +1,102 @@
+"""
+Generate metric deletion jobs and schedules.
+"""
+
+import os
+
+from dagster import (
+    MAX_RUNTIME_SECONDS_TAG,
+    DefaultScheduleStatus,
+    JobDefinition,
+    ScheduleDefinition,
+    get_dagster_logger,
+    job,
+    op,
+)
+
+from anomstack.config import specs
+from anomstack.jinja.render import render
+from anomstack.sql.read import read_sql
+
+ANOMSTACK_MAX_RUNTIME_SECONDS_TAG = os.getenv(
+    "ANOMSTACK_MAX_RUNTIME_SECONDS_TAG", 3600
+)
+
+
+def build_delete_job(spec: dict) -> JobDefinition:
+    """
+    Build job definitions for delete jobs.
+
+    Args:
+        spec (dict): A dictionary containing the specifications for the delete job.
+
+    Returns:
+        JobDefinition: A job definition for the delete job.
+    """
+
+    if spec.get("disable_delete"):
+
+        @job(
+            name=f'{spec["metric_batch"]}_delete_disabled',
+            tags={MAX_RUNTIME_SECONDS_TAG: ANOMSTACK_MAX_RUNTIME_SECONDS_TAG},
+        )
+        def _dummy_job():
+            @op(name=f'{spec["metric_batch"]}_delete_noop')
+            def noop():
+                pass
+
+            noop()
+
+        return _dummy_job
+
+    get_dagster_logger()
+
+    metric_batch = spec["metric_batch"]
+    db = spec["db"]
+    spec["alert_methods"]
+    spec["table_key"]
+    spec.get("metric_tags", {})
+    spec.get("delete_after_n_days", None)
+
+    @job(
+        name=f"{metric_batch}_delete",
+        tags={MAX_RUNTIME_SECONDS_TAG: ANOMSTACK_MAX_RUNTIME_SECONDS_TAG},
+    )
+    def _job():
+        """
+        Run delete job.
+        """
+
+        @op(name=f"{metric_batch}_delete_old_data")
+        def delete_old_data() -> None:
+            """
+            Delete old data.
+
+            Returns:
+                None: None.
+            """
+            _ = read_sql(render("delete_sql", spec), db, returns_df=False)
+
+            return None
+
+        delete_old_data()
+
+    return _job
+
+
+# Build delete jobs and schedules.
+delete_jobs = []
+delete_schedules = []
+for spec_name, spec in specs.items():
+    delete_job = build_delete_job(spec)
+    delete_jobs.append(delete_job)
+    if spec["delete_default_schedule_status"] == "RUNNING":
+        delete_default_schedule_status = DefaultScheduleStatus.RUNNING
+    else:
+        delete_default_schedule_status = DefaultScheduleStatus.STOPPED
+    delete_schedule = ScheduleDefinition(
+        job=delete_job,
+        cron_schedule=spec["delete_cron_schedule"],
+        default_status=delete_default_schedule_status,
+    )
+    delete_schedules.append(delete_schedule)

anomstack/main.py

@@ -6,6 +6,7 @@
 
 from anomstack.jobs.alert import alert_jobs, alert_schedules
 from anomstack.jobs.change import change_jobs, change_schedules
+from anomstack.jobs.delete import delete_jobs, delete_schedules
 from anomstack.jobs.ingest import ingest_jobs, ingest_schedules
 from anomstack.jobs.llmalert import llmalert_jobs, llmalert_schedules
 from anomstack.jobs.plot import plot_jobs, plot_schedules
@@ -23,6 +24,7 @@
     + plot_jobs
     + change_jobs
     + summary_jobs
+    + delete_jobs
 )
 sensors = [email_on_run_failure]
 schedules = (
@@ -34,6 +36,7 @@
     + plot_schedules
     + change_schedules
     + summary_schedules
+    + delete_schedules
 )
 
 defs = Definitions(

anomstack/sql/read.py

@@ -3,55 +3,28 @@
 different database connectors.
 """
 
-import re
 
 import pandas as pd
-import sqlglot
 from dagster import get_dagster_logger
 
 from anomstack.df.utils import log_df_info
-from anomstack.external.duckdb.duckdb import read_sql_duckdb
+from anomstack.external.duckdb.duckdb import read_sql_duckdb, run_sql_duckdb
 from anomstack.external.gcp.bigquery import read_sql_bigquery
 from anomstack.external.snowflake.snowflake import read_sql_snowflake
-from anomstack.external.sqlite.sqlite import read_sql_sqlite
+from anomstack.external.sqlite.sqlite import read_sql_sqlite, run_sql_sqlite
+from anomstack.sql.translate import db_translate
 
 pd.options.display.max_columns = 10
 
 
-def db_translate(sql: str, db: str) -> str:
-    """
-    Replace some functions with their db-specific equivalents.
-
-    Args:
-        sql (str): The SQL query to be translated.
-        db (str): The name of the database to which the query will be sent.
-
-    Returns:
-        str: The translated SQL query.
-    """
-    # Transpile the SQL query to the target database dialect
-    sql = sqlglot.transpile(sql, write=db, identify=True, pretty=True)[0]
-    # Replace some functions with their db-specific equivalents
-    if db == "sqlite":
-        sql = sql.replace("GET_CURRENT_TIMESTAMP()", "DATETIME('now')")
-    elif db == "bigquery":
-        sql = sql.replace("GET_CURRENT_TIMESTAMP()", "CURRENT_TIMESTAMP()")
-        sql = re.sub(
-            r"DATE\('now', '(-?\d+) day'\)",
-            "DATE_ADD(CURRENT_DATE(), INTERVAL \\1 DAY)",
-            sql
-        )
-
-    return sql
-
-
-def read_sql(sql: str, db: str) -> pd.DataFrame:
+def read_sql(sql: str, db: str, returns_df: bool = True) -> pd.DataFrame:
     """
     Read data from SQL.
 
     Args:
         sql (str): SQL query to execute.
         db (str): Name of the database to connect to.
+        returns_df (bool, optional): Whether the query expects a DataFrame as a result.
 
     Returns:
         pd.DataFrame: A pandas DataFrame containing the results of the SQL query.
@@ -64,13 +37,31 @@ def read_sql(sql: str, db: str) -> pd.DataFrame:
     logger.debug(f"-- read_sql() is about to read this qry:\n{sql}")
 
     if db == "bigquery":
-        df = read_sql_bigquery(sql)
+        if returns_df:
+            df = read_sql_bigquery(sql)
+        elif not returns_df:
+            raise NotImplementedError(
+                "BigQuery not yet implemented for non-returns_df queries."
+            )
     elif db == "snowflake":
-        df = read_sql_snowflake(sql)
+        if returns_df:
+            df = read_sql_snowflake(sql)
+        elif not returns_df:
+            raise NotImplementedError(
+                "Snowflake not yet implemented for non-returns_df queries."
+            )
     elif db == "duckdb":
-        df = read_sql_duckdb(sql)
+        if returns_df:
+            df = read_sql_duckdb(sql)
+        elif not returns_df:
+            run_sql_duckdb(sql)
+            df = pd.DataFrame()
     elif db == "sqlite":
-        df = read_sql_sqlite(sql)
+        if returns_df:
+            df = read_sql_sqlite(sql)
+        elif not returns_df:
+            run_sql_sqlite(sql)
+            df = pd.DataFrame()
     else:
         raise ValueError(f"Unknown db: {db}")
 

anomstack/sql/translate.py

@@ -0,0 +1,30 @@
+import re
+
+import sqlglot
+
+
+def db_translate(sql: str, db: str) -> str:
+    """
+    Replace some functions with their db-specific equivalents.
+
+    Args:
+        sql (str): The SQL query to be translated.
+        db (str): The name of the database to which the query will be sent.
+
+    Returns:
+        str: The translated SQL query.
+    """
+    # Transpile the SQL query to the target database dialect
+    sql = sqlglot.transpile(sql, write=db, identify=True, pretty=True)[0]
+    # Replace some functions with their db-specific equivalents
+    if db == "sqlite":
+        sql = sql.replace("GET_CURRENT_TIMESTAMP()", "DATETIME('now')")
+    elif db == "bigquery":
+        sql = sql.replace("GET_CURRENT_TIMESTAMP()", "CURRENT_TIMESTAMP()")
+        sql = re.sub(
+            r"DATE\('now', '(-?\d+) day'\)",
+            "DATE_ADD(CURRENT_DATE(), INTERVAL \\1 DAY)",
+            sql
+        )
+
+    return sql

metrics/defaults/defaults.yaml

@@ -87,6 +87,8 @@ alert_cron_schedule: "*/5 * * * *" # cron schedule for alerting jobs
 alert_default_schedule_status: 'STOPPED' # default schedule status for alert jobs (RUNNING or STOPPED)
 change_cron_schedule: "*/5 * * * *" # cron schedule for change detection jobs
 change_default_schedule_status: 'STOPPED' # default schedule status for alert jobs (RUNNING or STOPPED)
+delete_cron_schedule: "0 6 * * *" # cron schedule for delete jobs
+delete_default_schedule_status: 'STOPPED' # default schedule status for delete jobs (RUNNING or STOPPED)
 llmalert_cron_schedule: "*/5 * * * *" # cron schedule for llmalerting jobs
 llmalert_default_schedule_status: 'STOPPED' # default schedule status for llmalert jobs (RUNNING or STOPPED)
 plot_cron_schedule: "*/5 * * * *" # cron schedule for plot jobs
@@ -116,6 +118,9 @@ llmalert_sql: >
 # default templated dashboard sql
 dashboard_sql: >
   {% include "./defaults/sql/dashboard.sql" %}
+# default templated delete sql
+delete_sql: >
+  {% include "./defaults/sql/delete.sql" %}
 # default templated summary sql
 summary_sql: >
   {% include "./defaults/sql/summary.sql" %}
@@ -143,3 +148,5 @@ disable_change: False # if you want to disable change detection job for some rea
 disable_llmalert: True # if you want to disable llmalert job for some reason.
 disable_plot: False # if you want to disable plot job for some reason.
 disable_summary: False # if you want to disable summary job for some reason.
+disable_delete: False # if you want to disable delete job for some reason.
+delete_after_n_days: 365 # delete metrics older than this number of days.

metrics/defaults/sql/delete.sql

@@ -0,0 +1,13 @@
+/*
+Template for generating the input data for the delete job.
+
+Written for SQLite but will be translated to target dialect based on `db` param via sqlglot.
+*/
+
+delete from {{ table_key }}
+where
+  metric_batch = '{{ metric_batch }}'
+  and
+  -- Delete data older than the specified number of days
+  date(metric_timestamp) < date('now', '-{{ delete_after_n_days }} day')
+;

tests/test_main.py

@@ -8,19 +8,19 @@
 
 
 def test_jobs_len():
-    assert len(jobs) == 127
+    assert len(jobs) == 145
 
 
 def test_jobs_len_ingest():
-    assert len(ingest_jobs) == (len(jobs)-1) / 7
+    assert len(ingest_jobs) == (len(jobs)-1) / 8
 
 
 def test_schedules_len():
-    assert len(schedules) == 127
+    assert len(schedules) == 145
 
 
 def test_schedules_len_ingest():
-    assert len(ingest_schedules) == (len(schedules)-1) / 7
+    assert len(ingest_schedules) == (len(schedules)-1) / 8
 
 
 def test_jobs_schedules_len_match():