Dryrun caching

Author: scholtzan

bigquery_etl/cli/query.py

@@ -2300,6 +2300,9 @@ def _update_query_schema(
         query_schema = Schema.from_query_file(
             query_file_path,
             content=sql_content,
+            project=project_name,
+            dataset=dataset_name,
+            table=table_name,
             use_cloud_function=use_cloud_function,
             respect_skip=respect_dryrun_skip,
             sql_dir=sql_dir,

bigquery_etl/dryrun.py

@@ -12,10 +12,14 @@
 """
 
 import glob
+import hashlib
 import json
+import os
+import pickle
 import random
 import re
 import sys
+import tempfile
 import time
 from enum import Enum
 from os.path import basename, dirname, exists
@@ -106,12 +110,12 @@ def __init__(
         dataset=None,
         table=None,
         billing_project=None,
-        ignore_content=False,
+        use_cache=True,
     ):
         """Instantiate DryRun class."""
         self.sqlfile = sqlfile
         self.content = content
-        self.ignore_content = ignore_content
+        self.use_cache = use_cache
         self.query_parameters = query_parameters
         self.strip_dml = strip_dml
         self.use_cloud_function = use_cloud_function
@@ -227,16 +231,125 @@ def get_sql(self):
 
         return sql
 
+    def _get_cache_key(self, sql):
+        """Generate cache key based on SQL content and other parameters."""
+        cache_input = f"{sql}|{self.project}|{self.dataset}|{self.table}"
+        return hashlib.sha256(cache_input.encode()).hexdigest()
+
+    def _get_cached_result(self, cache_key, ttl_seconds=None):
+        """Load cached dry run result from disk."""
+        if ttl_seconds is None:
+            ttl_seconds = ConfigLoader.get("dry_run", "cache_ttl_seconds", fallback=900)
+
+        cache_dir = os.path.join(tempfile.gettempdir(), "bigquery_etl_dryrun_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+        cache_file = os.path.join(cache_dir, f"dryrun_{cache_key}.pkl")
+
+        if os.path.exists(cache_file):
+            # check if cache is expired
+            file_age = time.time() - os.path.getmtime(cache_file)
+            if file_age > ttl_seconds:
+                try:
+                    os.remove(cache_file)
+                except OSError:
+                    pass
+                return None
+
+            try:
+                with open(cache_file, "rb") as f:
+                    cached_data = pickle.load(f)
+                cache_age = time.time() - os.path.getmtime(cache_file)
+                print(f"[DRYRUN CACHE HIT] {self.sqlfile} (age: {cache_age:.0f}s)")
+                return cached_data
+            except (pickle.PickleError, EOFError, OSError) as e:
+                print(f"[DRYRUN CACHE] Failed to load cache: {e}")
+                return None
+
+        return None
+
+    def _save_cached_result(self, cache_key, result):
+        """Save dry run result to disk cache."""
+        cache_dir = os.path.join(tempfile.gettempdir(), "bigquery_etl_dryrun_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+        cache_file = os.path.join(cache_dir, f"dryrun_{cache_key}.pkl")
+
+        try:
+            with open(cache_file, "wb") as f:
+                pickle.dump(result, f)
+
+            # save table metadata separately if present
+            if (
+                result
+                and "tableMetadata" in result
+                and self.project
+                and self.dataset
+                and self.table
+            ):
+                table_identifier = f"{self.project}.{self.dataset}.{self.table}"
+                self._save_cached_table_metadata(
+                    table_identifier, result["tableMetadata"]
+                )
+        except (pickle.PickleError, OSError) as e:
+            print(f"[DRYRUN CACHE] Failed to save cache: {e}")
+
+    def _get_cached_table_metadata(self, table_identifier, ttl_seconds=None):
+        """Load cached table metadata from disk based on table identifier."""
+        if ttl_seconds is None:
+            ttl_seconds = ConfigLoader.get("dry_run", "cache_ttl_seconds", fallback=900)
+
+        cache_dir = os.path.join(tempfile.gettempdir(), "bigquery_etl_dryrun_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+        # table identifier as cache key
+        table_cache_key = hashlib.sha256(table_identifier.encode()).hexdigest()
+        cache_file = os.path.join(cache_dir, f"table_metadata_{table_cache_key}.pkl")
+
+        if os.path.exists(cache_file):
+            # check if cache is expired
+            file_age = time.time() - os.path.getmtime(cache_file)
+
+            if file_age > ttl_seconds:
+                try:
+                    os.remove(cache_file)
+                except OSError:
+                    pass
+                return None
+
+            try:
+                with open(cache_file, "rb") as f:
+                    cached_data = pickle.load(f)
+                return cached_data
+            except (pickle.PickleError, EOFError, OSError) as e:
+                return None
+        return None
+
+    def _save_cached_table_metadata(self, table_identifier, metadata):
+        """Save table metadata to disk cache."""
+        cache_dir = os.path.join(tempfile.gettempdir(), "bigquery_etl_dryrun_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+        table_cache_key = hashlib.sha256(table_identifier.encode()).hexdigest()
+        cache_file = os.path.join(cache_dir, f"table_metadata_{table_cache_key}.pkl")
+
+        try:
+            with open(cache_file, "wb") as f:
+                pickle.dump(metadata, f)
+        except (pickle.PickleError, OSError) as e:
+            print(f"[TABLE METADATA] Failed to save cache for {table_identifier}: {e}")
+
     @cached_property
     def dry_run_result(self):
         """Dry run the provided SQL file."""
-        if self.ignore_content:
-            sql = None
+        if self.content:
+            sql = self.content
         else:
-            if self.content:
-                sql = self.content
-            elif self.content != "":
-                sql = self.get_sql()
+            sql = self.get_sql()
+
+        # Check cache first (if caching is enabled)
+        if sql is not None and self.use_cache:
+            cache_key = self._get_cache_key(sql)
+            cached_result = self._get_cached_result(cache_key)
+            if cached_result is not None:
+                self.dry_run_duration = 0  # Cached result, no actual dry run
+                return cached_result
 
         query_parameters = []
         if self.query_parameters:
@@ -356,6 +469,11 @@ def dry_run_result(self):
                     }
 
             self.dry_run_duration = time.time() - start_time
+
+            # Save to cache (if caching is enabled)
+            if self.use_cache:
+                self._save_cached_result(cache_key, result)
+
             return result
 
         except Exception as e:
@@ -481,6 +599,13 @@ def get_table_schema(self):
         ):
             return self.dry_run_result["tableMetadata"]["schema"]
 
+        # Check if table metadata is cached (if caching is enabled)
+        if self.use_cache and self.project and self.dataset and self.table:
+            table_identifier = f"{self.project}.{self.dataset}.{self.table}"
+            cached_metadata = self._get_cached_table_metadata(table_identifier)
+            if cached_metadata:
+                return cached_metadata["schema"]
+
         return []
 
     def get_dataset_labels(self):

bigquery_etl/schema/__init__.py

@@ -78,8 +78,6 @@ def for_table(
                     project=project,
                     dataset=dataset,
                     table=table,
-                    respect_skip=False,
-                    ignore_content=True,
                     *args,
                     **kwargs,
                 ).get_table_schema()

bigquery_etl/schema/stable_table_schema.py

@@ -3,6 +3,7 @@
 import json
 import os
 import pickle
+import shutil
 import tarfile
 import tempfile
 import urllib.request
@@ -51,6 +52,17 @@ def sortkey(self):
         )
 
 
+def _clear_dryrun_cache():
+    """Clear dry run cache when new schemas are downloaded."""
+    cache_dir = os.path.join(tempfile.gettempdir(), "bigquery_etl_dryrun_cache")
+    if os.path.exists(cache_dir):
+        try:
+            shutil.rmtree(cache_dir)
+            print(f"Cleared dry run cache at {cache_dir}")
+        except OSError as e:
+            print(f"Warning: Failed to clear dry run cache: {e}")
+
+
 def prod_schemas_uri():
     """Return URI for the schemas tarball deployed to shared-prod.
 
@@ -59,7 +71,7 @@ def prod_schemas_uri():
     with the most recent production schemas deploy.
     """
     dryrun = DryRun(
-        "moz-fx-data-shared-prod/telemetry_derived/foo/query.sql", content="SELECT 1"
+        "moz-fx-data-shared-prod/telemetry_derived/foo/query.sql", content="SELECT 1", use_cache=False
     )
     build_id = dryrun.get_dataset_labels()["schemas_build_id"]
     commit_hash = build_id.split("_")[-1]
@@ -88,6 +100,11 @@ def get_stable_table_schemas() -> List[SchemaFile]:
             print(f"Failed to load cached schemas: {e}, re-downloading...")
 
     print(f"Downloading schemas from {schemas_uri}")
+
+    # Clear dry run cache when downloading new schemas
+    # Schema changes could affect dry run results
+    _clear_dryrun_cache()
+
     with urllib.request.urlopen(schemas_uri) as f:
         tarbytes = BytesIO(f.read())
 

bqetl_project.yaml

@@ -32,6 +32,7 @@ dry_run:
   function_accounts:
   - bigquery-etl-dryrun@moz-fx-data-shared-prod.iam.gserviceaccount.com
   - bigquery-etl-dryrun@moz-fx-data-shar-nonprod-efed.iam.gserviceaccount.com
+  cache_ttl_seconds: 900  # Cache dry run results for 15 minutes (900 seconds)
   skip:
   ## skip all data-observability-dev queries due to CI lacking permissions in that project.
   # TODO: once data observability platform assessment concludes this should be removed.