Merge pull request #6 from openstates/set-prefix-bucket-directory

Use DAG run start time

Author: alexobaseki

.gitignore

@@ -2,6 +2,6 @@ logs
 _data
 .cache
 merged_*
-db.db
+db.*
 __pycache__
 *json

audits/event.sql

@@ -1,4 +1,4 @@
--- all events have sponsors?
+-- all events are classified?
 AUDIT (
     name assert_events_are_classified,
     blocking false

main.py

@@ -1,8 +1,11 @@
 import argparse
 
 from sqlmesh_tasks import sqlmesh_plan
+from openstates_metadata import lookup
+from utils import send_slack_message
 
-if __name__ == "__main__":
+
+def main() -> None:
     default_parser = argparse.ArgumentParser(add_help=False)
 
     parser = argparse.ArgumentParser(
@@ -25,10 +28,26 @@
 
     args = parser.parse_args()
     entity = args.entity
-    jurisdiction = args.jurisdiction
+    jur_obj = lookup(abbr=args.jurisdiction) if args.jurisdiction else None
 
     if entity:
         entities = [entity]
     else:
         entities = ["bill", "event", "vote_event"]
-    report = sqlmesh_plan(entities, jurisdiction)
+
+    # Use Jurisdiction if it is provided
+    if jur_obj:
+        reports = sqlmesh_plan(entities, jur_obj.jurisdiction_id)
+    else:
+        reports = sqlmesh_plan(entities)
+
+    # Send report
+    if reports:
+        reports = "\n".join(reports)
+        jur_name = jur_obj.name if jur_obj else ""
+        msg = f"Scrape Output Audit for {jur_name}: \n{reports}"
+        send_slack_message("data-reports", msg)
+
+
+if __name__ == "__main__":
+    main()

poetry.lock

@@ -9,6 +9,8 @@ python = "^3.9"
 sqlmesh = "^0.176.0"
 duckdb = "^1.2.2"
 google-cloud-storage = "^3.1.0"
+slack-sdk = "^3.35.0"
+openstates-metadata = "^2024.10.3"
 
 [tool.poetry.group.dev.dependencies]
 black = "^22"

pyproject.toml

@@ -20,10 +20,10 @@ def extract_audit_error(stdout: str) -> typing.Union[str, None]:
     return None
 
 
-def sqlmesh_plan(entities: list[str], jurisdiction: str) -> list:
+def sqlmesh_plan(entities: list[str], jurisdiction: str = None) -> list:
     """Run SQLMesh plan on initialized DuckDB data"""
 
-    initialize_entities = init_duckdb(jurisdiction, entities)
+    initialize_entities = init_duckdb(entities, jurisdiction)
     initialize_entities = [f"staged.{entity}" for entity in initialize_entities]
     reports = []
 
@@ -60,7 +60,7 @@ def sqlmesh_plan(entities: list[str], jurisdiction: str) -> list:
             raise
 
         if report:
-            logger.info(f"Entity: {entity} audit failed:\n", report)
+            logger.info(f"Entity: {entity} audit failed {report}")
             reports.append(report)
         else:
             logger.info(f"Entity: {entity} audit passed.")

sqlmesh_tasks.py

@@ -4,13 +4,36 @@
 
 import duckdb
 from google.cloud import storage
+from slack_sdk.web import WebClient
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("openstates")
 
 GCP_PROJECT = os.environ.get("GCP_PROJECT", None)
 BUCKET_NAME = os.environ.get("BUCKET_NAME", None)
 SCRAPE_LAKE_PREFIX = os.environ.get("BUCKET_PREFIX", "legislation")
+DAG_RUN_START = os.environ.get("DAG_RUN_START", None)
+SLACK_BEARER_TOKEN = os.environ.get("SLACK_BEARER_TOKEN", None)
+
+
+def send_slack_message(channel, msg=None, attachments=None) -> None:
+    if not SLACK_BEARER_TOKEN:
+        logger.warning("No SLACK_BEARER_TOKEN, cannot send slack notification.")
+        return
+    sendobj = {"channel": channel}
+    if msg:
+        sendobj["text"] = msg
+    if attachments:
+        if len(attachments) > 50:
+            attachments.insert(0, {"title": "Too many attachments", "color": "FF3333"})
+            attachments = attachments[:50]
+        sendobj["attachments"] = attachments
+    if sendobj.get("text", "") or sendobj.get("attachments", ""):
+        try:
+            client = WebClient(token=SLACK_BEARER_TOKEN)
+            client.chat_postMessage(**sendobj)
+        except Exception as e:
+            logger.error(f"Couldn't send slack message: {e}")
 
 
 def check_for_json_files(file_path: str) -> bool:
@@ -50,35 +73,40 @@ def download_files_from_gcs(file_path: str) -> None:
 
 
 def init_duckdb(
-    jurisdiction: str,
     entities: list[str],
-    last_scrape_end_time: str = None,
+    jurisdiction: str = None,
 ) -> list[str]:
     """Initialize Duckdb and load data, return list of tables created for usage downstream."""
 
     db_path = "db.db"
     if os.path.exists(db_path):
         os.remove(db_path)
 
-    sub_directory = "*"
-    if jurisdiction and last_scrape_end_time:
-        sub_directory = jurisdiction.replace("ocd-jurisdiction/", "")
-        sub_directory = f"{sub_directory}/{last_scrape_end_time}"
+    # Determine subdirectory pattern for file search
+    if jurisdiction and DAG_RUN_START:
+        # Strip OCD prefix and build a dynamic path using DAG_RUN_START
+        relative_path = jurisdiction.replace("ocd-jurisdiction/", "")
+        sub_directory = f"**/{relative_path}/{DAG_RUN_START}"
+    else:
+        sub_directory = "**"
+
     # Create DuckDB and load
     logger.info("Creating DuckDB schema and loading data...")
     con = duckdb.connect(db_path)
     con.execute("CREATE SCHEMA IF NOT EXISTS scraper")
     table_created = []
 
-    file_path_prefix = f"./*/{sub_directory}"
+    file_path_prefix = f"./{sub_directory}"
     all_files_path = f"{file_path_prefix}/*.json"
 
     # Grab scrape output from data lake if none is provided
     if not check_for_json_files(all_files_path):
         logger.info(
             "No file found in local directory, attempting to download from GCS, requires credentials in ENV."
         )
-        download_files_from_gcs(sub_directory)
+        # Remove "**/" from path prefix before passing to GCS downloader
+        gcs_path = sub_directory[3:] if sub_directory.startswith("**/") else sub_directory
+        download_files_from_gcs(gcs_path)
 
     # Load data into duckdb table
     for entity in entities: