Fix for_table to use caching

Author: scholtzan

bigquery_etl/dryrun.py

@@ -18,6 +18,7 @@
 import pickle
 import random
 import re
+import shutil
 import sys
 import tempfile
 import time
@@ -198,6 +199,17 @@ def skipped_files(sql_dir=ConfigLoader.get("default", "sql_dir")) -> Set[str]:
 
         return skip_files
 
+    @staticmethod
+    def clear_cache():
+        """Clear dry run cache directory."""
+        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 skip(self):
         """Determine if dry run should be skipped."""
         return self.respect_skip and self.sqlfile in self.skipped_files(
@@ -470,8 +482,9 @@ def dry_run_result(self):
 
             self.dry_run_duration = time.time() - start_time
 
-            # Save to cache (if caching is enabled)
-            if self.use_cache:
+            # Save to cache (if caching is enabled and result is valid)
+            # Don't cache errors to allow retries
+            if self.use_cache and result.get("valid"):
                 self._save_cached_result(cache_key, result)
 
             return result

bigquery_etl/schema/__init__.py

@@ -13,6 +13,7 @@
 from google.cloud.bigquery import SchemaField
 
 from .. import dryrun
+from ..config import ConfigLoader
 
 SCHEMA_FILE = "schema.yaml"
 
@@ -58,24 +59,34 @@ def from_json(cls, json_schema):
         return cls(json_schema)
 
     @classmethod
-    def for_table(cls, project, dataset, table, partitioned_by=None, *args, **kwargs):
+    def for_table(
+        cls,
+        project,
+        dataset,
+        table,
+        partitioned_by=None,
+        filename="query.sql",
+        *args,
+        **kwargs,
+    ):
         """Get the schema for a BigQuery table."""
         query = f"SELECT * FROM `{project}.{dataset}.{table}`"
 
         if partitioned_by:
             query += f" WHERE DATE(`{partitioned_by}`) = DATE('2020-01-01')"
 
         try:
+            sql_dir = ConfigLoader.get("default", "sql_dir")
             return cls(
                 dryrun.DryRun(
-                    os.path.join(project, dataset, table, "query.sql"),
+                    os.path.join(sql_dir, project, dataset, table, filename),
                     query,
                     project=project,
                     dataset=dataset,
                     table=table,
                     *args,
                     **kwargs,
-                ).get_schema()
+                ).get_table_schema()
             )
         except Exception as e:
             print(f"Cannot get schema for {project}.{dataset}.{table}: {e}")

bigquery_etl/schema/stable_table_schema.py

@@ -3,7 +3,6 @@
 import json
 import os
 import pickle
-import shutil
 import tarfile
 import tempfile
 import urllib.request
@@ -52,17 +51,6 @@ 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.
 
@@ -105,7 +93,7 @@ def get_stable_table_schemas() -> List[SchemaFile]:
 
     # Clear dry run cache when downloading new schemas
     # Schema changes could affect dry run results
-    _clear_dryrun_cache()
+    DryRun.clear_cache()
 
     with urllib.request.urlopen(schemas_uri) as f:
         tarbytes = BytesIO(f.read())