[run-tests] undo schema changes

Author: scholtzan

bigquery_etl/dryrun.py

@@ -318,7 +318,7 @@ def _get_cached_table_metadata(self, table_identifier, ttl_seconds=None):
                 with open(cache_file, "rb") as f:
                     cached_data = pickle.load(f)
                 return cached_data
-            except (pickle.PickleError, EOFError, OSError) as e:
+            except (pickle.PickleError, EOFError, OSError):
                 return None
         return None
 

bigquery_etl/schema/__init__.py

@@ -13,7 +13,6 @@
 from google.cloud.bigquery import SchemaField
 
 from .. import dryrun
-from ..config import ConfigLoader
 
 SCHEMA_FILE = "schema.yaml"
 
@@ -59,28 +58,24 @@ def from_json(cls, json_schema):
         return cls(json_schema)
 
     @classmethod
-    def for_table(
-        cls,
-        project,
-        dataset,
-        table,
-        partitioned_by=None,
-        filename="query.sql",
-        *args,
-        **kwargs,
-    ):
+    def for_table(cls, project, dataset, table, partitioned_by=None, *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(sql_dir, project, dataset, table, filename),
+                    os.path.join(project, dataset, table, "query.sql"),
+                    query,
                     project=project,
                     dataset=dataset,
                     table=table,
                     *args,
                     **kwargs,
-                ).get_table_schema()
+                ).get_schema()
             )
         except Exception as e:
             print(f"Cannot get schema for {project}.{dataset}.{table}: {e}")

bigquery_etl/schema/stable_table_schema.py

@@ -71,7 +71,9 @@ 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", use_cache=False
+        "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]

tests/test_dryrun.py

@@ -1,7 +1,4 @@
 import os
-import tempfile
-import time
-from unittest import mock
 
 import pytest