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 = Path(tempfile.gettempdir()) / "bigquery_etl_dryrun_cache"
+        if cache_dir.exists():
+            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(
@@ -241,41 +253,52 @@ def _get_cached_result(self, cache_key, ttl_seconds=None):
         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")
+        cache_dir = Path(tempfile.gettempdir()) / "bigquery_etl_dryrun_cache"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_file = cache_dir / f"dryrun_{cache_key}.pkl"
+
+        try:
+            if not cache_file.exists():
+                return None
 
-        if os.path.exists(cache_file):
             # check if cache is expired
-            file_age = time.time() - os.path.getmtime(cache_file)
+            file_age = time.time() - cache_file.stat().st_mtime
             if file_age > ttl_seconds:
                 try:
-                    os.remove(cache_file)
+                    cache_file.unlink()
                 except OSError:
                     pass
                 return None
 
+            cached_data = pickle.loads(cache_file.read_bytes())
+            cache_age = time.time() - cache_file.stat().st_mtime
+            print(f"[DRYRUN CACHE HIT] {self.sqlfile} (age: {cache_age:.0f}s)")
+            return cached_data
+        except (pickle.PickleError, EOFError, OSError, FileNotFoundError) as e:
+            print(f"[DRYRUN CACHE] Failed to load cache: {e}")
             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
+                if cache_file.exists():
+                    cache_file.unlink()
+            except OSError:
+                pass
+            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")
+        """Save dry run result to disk cache using atomic write."""
+        cache_dir = Path(tempfile.gettempdir()) / "bigquery_etl_dryrun_cache"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_file = cache_dir / f"dryrun_{cache_key}.pkl"
 
         try:
-            with open(cache_file, "wb") as f:
+            # write to temporary file first, then atomically rename
+            # this prevents race conditions where readers get partial files
+            temp_file = Path(str(cache_file) + f".tmp.{os.getpid()}")
+            with open(temp_file, "wb") as f:
                 pickle.dump(result, f)
+                f.flush()
+                os.fsync(f.fileno())  # Ensure data is written to disk
+
+            temp_file.replace(cache_file)
 
             # save table metadata separately if present
             if (
@@ -291,49 +314,73 @@ def _save_cached_result(self, cache_key, result):
                 )
         except (pickle.PickleError, OSError) as e:
             print(f"[DRYRUN CACHE] Failed to save cache: {e}")
+            try:
+                temp_file = Path(str(cache_file) + f".tmp.{os.getpid()}")
+                if temp_file.exists():
+                    temp_file.unlink()
+            except OSError:
+                pass
 
     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)
+        cache_dir = Path(tempfile.gettempdir()) / "bigquery_etl_dryrun_cache"
+        cache_dir.mkdir(parents=True, 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")
+        cache_file = cache_dir / f"table_metadata_{table_cache_key}.pkl"
+
+        try:
+            if not cache_file.exists():
+                return None
 
-        if os.path.exists(cache_file):
             # check if cache is expired
-            file_age = time.time() - os.path.getmtime(cache_file)
+            file_age = time.time() - cache_file.stat().st_mtime
 
             if file_age > ttl_seconds:
                 try:
-                    os.remove(cache_file)
+                    cache_file.unlink()
                 except OSError:
                     pass
                 return None
 
+            cached_data = pickle.loads(cache_file.read_bytes())
+            return cached_data
+        except (pickle.PickleError, EOFError, OSError, FileNotFoundError) as e:
+            print(f"[TABLE METADATA] Failed to load cache for {table_identifier}: {e}")
             try:
-                with open(cache_file, "rb") as f:
-                    cached_data = pickle.load(f)
-                return cached_data
-            except (pickle.PickleError, EOFError, OSError):
-                return None
-        return None
+                if cache_file.exists():
+                    cache_file.unlink()
+            except OSError:
+                pass
+            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)
+        """Save table metadata to disk cache using atomic write."""
+        cache_dir = Path(tempfile.gettempdir()) / "bigquery_etl_dryrun_cache"
+        cache_dir.mkdir(parents=True, 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")
+        cache_file = cache_dir / f"table_metadata_{table_cache_key}.pkl"
 
         try:
-            with open(cache_file, "wb") as f:
+            # write to temporary file first, then atomically rename
+            temp_file = Path(str(cache_file) + f".tmp.{os.getpid()}")
+            with open(temp_file, "wb") as f:
                 pickle.dump(metadata, f)
+                f.flush()
+                os.fsync(f.fileno())
+
+            temp_file.replace(cache_file)
         except (pickle.PickleError, OSError) as e:
             print(f"[TABLE METADATA] Failed to save cache for {table_identifier}: {e}")
+            try:
+                temp_file = Path(str(cache_file) + f".tmp.{os.getpid()}")
+                if temp_file.exists():
+                    temp_file.unlink()
+            except OSError:
+                pass
 
     @cached_property
     def dry_run_result(self):
@@ -343,7 +390,7 @@ def dry_run_result(self):
         else:
             sql = self.get_sql()
 
-        # Check cache first (if caching is enabled)
+        # 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)
@@ -470,8 +517,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())

tests/test_dryrun.py

@@ -212,8 +212,11 @@ def test_cache_key_generation(self, tmp_query_path):
         cache_key3 = dryrun._get_cache_key(different_sql)
         assert cache_key1 != cache_key3
 
-    def test_cache_save_and_load(self, tmp_query_path):
+    def test_cache_save_and_load(self, tmp_query_path, monkeypatch, tmp_path):
         """Test that dry run results can be saved and loaded from cache."""
+        # Use isolated cache directory for this test to avoid interference from other tests
+        monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
+
         query_file = tmp_query_path / "query.sql"
         query_file.write_text("SELECT 123")
 
@@ -237,8 +240,11 @@ def test_cache_save_and_load(self, tmp_query_path):
         assert cached_result["valid"] is True
         assert cached_result["schema"]["fields"][0]["name"] == "test"
 
-    def test_cache_expiration(self, tmp_query_path):
+    def test_cache_expiration(self, tmp_query_path, monkeypatch, tmp_path):
         """Test that cache expires after TTL."""
+        # Use isolated cache directory for this test to avoid interference from other tests
+        monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
+
         query_file = tmp_query_path / "query.sql"
         query_file.write_text("SELECT 123")
 
@@ -257,8 +263,11 @@ def test_cache_expiration(self, tmp_query_path):
         expired = dryrun._get_cached_result(cache_key, ttl_seconds=0)
         assert expired is None
 
-    def test_cache_respects_sql_changes(self, tmp_query_path):
+    def test_cache_respects_sql_changes(self, tmp_query_path, monkeypatch, tmp_path):
         """Test that changing SQL content creates a different cache entry."""
+        # Use isolated cache directory for this test to avoid interference from other tests
+        monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
+
         query_file = tmp_query_path / "query.sql"
 
         # First SQL
@@ -286,8 +295,11 @@ def test_cache_respects_sql_changes(self, tmp_query_path):
         cached2 = dryrun2._get_cached_result(cache_key2)
         assert cached2 is None
 
-    def test_table_metadata_cache(self, tmp_query_path):
+    def test_table_metadata_cache(self, tmp_query_path, monkeypatch, tmp_path):
         """Test that table metadata can be cached by table identifier."""
+        # Use isolated cache directory for this test to avoid interference from other tests
+        monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
+
         query_file = tmp_query_path / "query.sql"
         query_file.write_text("SELECT 123")
 
@@ -314,8 +326,13 @@ def test_table_metadata_cache(self, tmp_query_path):
         assert cached_metadata["schema"]["fields"][0]["name"] == "col1"
         assert cached_metadata["tableType"] == "TABLE"
 
-    def test_table_metadata_cache_different_tables(self, tmp_query_path):
+    def test_table_metadata_cache_different_tables(
+        self, tmp_query_path, monkeypatch, tmp_path
+    ):
         """Test that different tables have separate cache entries."""
+        # Use isolated cache directory for this test to avoid interference from other tests
+        monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
+
         query_file = tmp_query_path / "query.sql"
         query_file.write_text("SELECT 123")