Enhance SQLite interaction by adding DataFrame return option and create index listing script

Author: andrewm4894

anomstack/external/sqlite/sqlite.py

@@ -188,22 +188,27 @@ def _action():
     return with_sqlite_retry(_action, logger=logger)
 
 
-def run_sql_sqlite(sql: str) -> None:
+def run_sql_sqlite(sql: str, return_df: bool = False):
     """
     Execute a non-returning SQL statement (e.g. CREATE, INSERT, UPDATE, DELETE) with retry logic.
 
     Args:
         sql (str): The SQL statement to execute.
+        return_df (bool, optional): Whether to return the result as a DataFrame. Defaults to False.
 
     Returns:
         None
     """
     logger = get_dagster_logger()
     logger.info(f"Executing SQL against DB path: {get_sqlite_path()}")
 
-    def _action():
+    def _action(return_df=return_df):
         with sqlite_connection() as conn:
-            conn.execute(sql)
+            cursor = conn.execute(sql)
             conn.commit()
+            rows = cursor.fetchall()
+            columns = [desc[0] for desc in cursor.description] if cursor.description else get_columns_from_sql(sql)
+            df = pd.DataFrame(rows, columns=columns)
+            return df if return_df else None
 
-    with_sqlite_retry(_action, logger=logger)
+    return with_sqlite_retry(_action, logger=logger)

scripts/sqlite/list_indexes.py

@@ -0,0 +1,14 @@
+from anomstack.external.sqlite.sqlite import run_sql_sqlite
+from dotenv import load_dotenv
+
+load_dotenv(override=True)
+
+
+def list_indexes():
+    query = "SELECT name, tbl_name FROM sqlite_master WHERE type='index' ORDER BY name;"
+    df = run_sql_sqlite(query, return_df=True)
+    print(df)
+
+if __name__ == "__main__":
+    df = list_indexes()
+

scripts/sqlite/list_tables.py

@@ -1,24 +1,13 @@
 from anomstack.external.sqlite.sqlite import run_sql_sqlite
 from dotenv import load_dotenv
 
-load_dotenv()
+load_dotenv(override=True)
 
 
-def list_all_tables():
-    # Query to get all table names from sqlite_master
-    tables_query = "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name;"
-    
-    # Execute the query
-    tables = run_sql_sqlite(tables_query)
-    
-    if not tables:
-        print("No tables found.")
-        return
+def list_tables():
+    query = "SELECT name, tbl_name FROM sqlite_master WHERE type='table' ORDER BY name;"
+    df = run_sql_sqlite(query, return_df=True)
+    print(df)
 
-    # Print the list of tables
-    print("Tables in the database:")
-    for table in tables:
-        print(f"- {table[0]}")  # Extract and print table name
-
-# Run the function
-list_all_tables()
+if __name__ == "__main__":
+    list_tables()