patch query text for e2e test

censoredplanet · Sep 6, 2023 · 68e416b · 68e416b
1 parent 092d942
commit 68e416b
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 15 deletions.
diff --git a/pipeline/manual_e2e_test.py b/pipeline/manual_e2e_test.py
@@ -440,16 +440,28 @@ def test_pipeline_e2e(self) -> None:
 
       client = cloud_bigquery.Client(
           project=firehook_resources.DEV_PROJECT_NAME)
-      # pylint: disable=protected-access
-      # Write derived table
-      run_queries._run_query(
-          client,
+
+      with open(
           'table/queries/merged_reduced_scans.sql',
-          firehook_resources.DEV_PROJECT_NAME,
-          BEAM_TEST_BASE_DATASET,
-          BEAM_TEST_BASE_DATASET,
-      )
-      # pylint: enable=protected-access
+          encoding='utf-8') as query_file:
+        # Replace the year restriction in the query with a predefined date
+        # so the query will continue to work with the e2e test data.
+        # This string must stay updated to match the query string.
+        query_text = query_file.read()
+        fixed_query_text = query_text.replace(
+            'SET earliest_date = DATE_SUB(CURRENT_DATE, INTERVAL 2 YEAR);',
+            'SET earliest_date = \'2018-01-01\';')
+
+        # pylint: disable=protected-access
+        # Write derived table
+        run_queries._run_query_text(
+            client,
+            fixed_query_text,
+            firehook_resources.DEV_PROJECT_NAME,
+            BEAM_TEST_BASE_DATASET,
+            BEAM_TEST_BASE_DATASET,
+        )
+        # pylint: enable=protected-access
 
       written_derived_rows = get_bq_rows(client, [derived_table_name])
       self.assertEqual(len(written_derived_rows), 53)

diff --git a/table/run_queries.py b/table/run_queries.py
@@ -37,14 +37,21 @@
 def _run_query(client: cloud_bigquery.Client, filepath: str, project_name: str,
                base_dataset: str,
                derived_dataset: str) -> cloud_bigquery.table.RowIterator:
-  with open(filepath, encoding='utf-8') as sql:
-    query = sql.read()
+  with open(filepath, encoding='utf-8') as sql_file:
+    query_text = sql_file.read()
+    return _run_query_text(client, query_text, project_name, base_dataset,
+                           derived_dataset)
 
-    query = query.replace(PROJECT_NAME_PLACEHOLDER, project_name)
-    query = query.replace(BASE_PLACEHOLDER, base_dataset)
-    query = query.replace(DERIVED_PLACEHOLDER, derived_dataset)
 
-    query_job = client.query(query)
+# This is used to patch the query for E2E testing
+def _run_query_text(client: cloud_bigquery.Client, query: str,
+                    project_name: str, base_dataset: str,
+                    derived_dataset: str) -> cloud_bigquery.table.RowIterator:
+  query = query.replace(PROJECT_NAME_PLACEHOLDER, project_name)
+  query = query.replace(BASE_PLACEHOLDER, base_dataset)
+  query = query.replace(DERIVED_PLACEHOLDER, derived_dataset)
+
+  query_job = client.query(query)
   return query_job.result()