fix(tests): Ensure concurrent test runs do not operate on the same BQ datasets

cli/index_run_e2e_test.ts

@@ -3,6 +3,8 @@ import * as fs from "fs-extra";
 import { dump as dumpYaml, load as loadYaml } from "js-yaml";
 import * as path from "path";
 
+import * as dfapi from "df/cli/api";
+import { BigQueryDbAdapter } from "df/cli/api/dbadapters/bigquery";
 import {
   alterWorkflowSettings,
   CREDENTIALS_PATH,
@@ -14,9 +16,12 @@ import {
 } from "df/cli/index_test_base";
 import { version } from "df/core/version";
 import { suite, test, writeDefinitionFile } from "df/testing";
+import { dropDatasets, dropStaleDatasets, makeUniqueSuffix } from "df/testing/bigquery";
 import { integrationSuite } from "df/testing/credentials";
 import { TmpDirFixture } from "df/testing/fixtures";
 
+const ON_SCHEMA_CHANGE_DATASET = "dataform_e2e_osc";
+
 integrationSuite("run e2e", ({ afterEach }) => {
   const tmpDirFixture = new TmpDirFixture(afterEach);
 
@@ -424,9 +429,22 @@ select 2
     });
   });
 
-  suite("onSchemaChange", ({ beforeEach }) => {
+  suite("onSchemaChange", ({ before, beforeEach, after }) => {
     let projectDir: string;
-    const uniqueDataset = `dataform_e2e_osc_${Math.random().toString(36).substring(7)}`;
+    const dbadapter = new BigQueryDbAdapter(dfapi.credentials.read(CREDENTIALS_PATH));
+    const uniqueDataset = `${ON_SCHEMA_CHANGE_DATASET}_${makeUniqueSuffix()}`;
+
+    before("drop stale datasets", async () => {
+      await dropStaleDatasets(dbadapter, {
+        project: INTEGRATION_TEST_PROJECT,
+        location: INTEGRATION_TEST_LOCATION,
+        prefixes: [ON_SCHEMA_CHANGE_DATASET],
+      });
+    });
+
+    after("drop datasets", async () => {
+      await dropDatasets(dbadapter, INTEGRATION_TEST_PROJECT, [uniqueDataset]);
+    });
 
     beforeEach("setup test project", async () => {
       projectDir = tmpDirFixture.createNewTmpDir();

cli/index_test_base.ts

@@ -124,8 +124,9 @@ export function alterWorkflowSettings(
 export async function setupJitProject(
   tmpDirFixture: TmpDirFixture,
   projectDir: string,
+  workflowSettingsOverrides?: Partial<dataform.IWorkflowSettings>,
 ): Promise<void> {
-  await setupProject(tmpDirFixture, projectDir);
+  await setupProject(tmpDirFixture, projectDir, workflowSettingsOverrides);
   writeDefinitionFile(
     projectDir,
     "jit_table.js",

cli/tests/jit/index_jit_advanced_test.ts

@@ -1,12 +1,39 @@
 import { expect } from "chai";
 
-import { CREDENTIALS_PATH, runCli, setupJitProject } from "df/cli/index_test_base";
+import * as dfapi from "df/cli/api";
+import { BigQueryDbAdapter } from "df/cli/api/dbadapters/bigquery";
+import {
+  CREDENTIALS_PATH,
+  INTEGRATION_TEST_LOCATION,
+  INTEGRATION_TEST_PROJECT,
+  runCli,
+  setupJitProject,
+} from "df/cli/index_test_base";
 import { test, writeDefinitionFile } from "df/testing";
+import { dropDatasets, dropStaleDatasets, makeUniqueSuffix } from "df/testing/bigquery";
 import { integrationSuite } from "df/testing/credentials";
 import { TmpDirFixture } from "df/testing/fixtures";
 
-integrationSuite("JiT support advanced", ({ afterEach }) => {
+const JIT_INCREMENTAL_DATASET = "dataform_e2e_jit_inc";
+
+integrationSuite("JiT support advanced", ({ afterEach, before, after }) => {
   const tmpDirFixture = new TmpDirFixture(afterEach);
+  const dbadapter = new BigQueryDbAdapter(dfapi.credentials.read(CREDENTIALS_PATH));
+  const uniqueSuffix = makeUniqueSuffix();
+
+  before("drop stale datasets", async () => {
+    await dropStaleDatasets(dbadapter, {
+      project: INTEGRATION_TEST_PROJECT,
+      location: INTEGRATION_TEST_LOCATION,
+      prefixes: [JIT_INCREMENTAL_DATASET],
+    });
+  });
+
+  after("drop datasets", async () => {
+    await dropDatasets(dbadapter, INTEGRATION_TEST_PROJECT, [
+      `${JIT_INCREMENTAL_DATASET}_${uniqueSuffix}`,
+    ]);
+  });
 
   test("JiT preOps and postOps support", async () => {
     const projectDir = tmpDirFixture.createNewTmpDir();
@@ -104,7 +131,9 @@ integrationSuite("JiT support advanced", ({ afterEach }) => {
     { name: "JiT incremental mode validation with consecutive runs", timeout: 60000 },
     async () => {
       const projectDir = tmpDirFixture.createNewTmpDir();
-      await setupJitProject(tmpDirFixture, projectDir);
+      await setupJitProject(tmpDirFixture, projectDir, {
+        defaultDataset: `${JIT_INCREMENTAL_DATASET}_${uniqueSuffix}`,
+      });
       writeDefinitionFile(
         projectDir,
         "inc_jit.js",

testing/BUILD

@@ -17,6 +17,7 @@ js_library(
 ts_library(
     name = "testing",
     srcs = [
+        "bigquery.ts",
         "child_process.ts",
         "credentials.ts",
         "fixtures.ts",

testing/bigquery.ts

@@ -0,0 +1,77 @@
+import { randomBytes } from "crypto";
+
+/** The minimal interface needed to run queries, satisfied by `BigQueryDbAdapter`. */
+export interface IQueryExecutor {
+  execute(statement: string): Promise<{ rows: any[] }>;
+}
+
+const UNIQUE_SUFFIX_BYTES = 4;
+const UNIQUE_SUFFIX_PATTERN = `[0-9a-f]{${UNIQUE_SUFFIX_BYTES * 2}}`;
+const DATASET_PREFIX_REGEX = /^[A-Za-z0-9_]+$/;
+
+const DEFAULT_STALE_DATASET_MIN_AGE_HOURS = 6;
+
+export function makeUniqueSuffix(): string {
+  return randomBytes(UNIQUE_SUFFIX_BYTES).toString("hex");
+}
+
+/** Drops the given datasets, including their contents. Best effort: failures are only logged. */
+export async function dropDatasets(
+  executor: IQueryExecutor,
+  project: string,
+  datasets: string[],
+): Promise<void> {
+  await Promise.all(
+    datasets.map(async (dataset) => {
+      try {
+        await executor.execute(`drop schema if exists \`${project}.${dataset}\` cascade`);
+      } catch (e) {
+        console.warn(`Failed to drop dataset ${dataset}: ${e?.message || e}`);
+      }
+    }),
+  );
+}
+
+export interface IDropStaleDatasetsOptions {
+  project: string;
+  location: string;
+  prefixes: string[];
+  minAgeHours?: number;
+}
+
+export async function dropStaleDatasets(
+  executor: IQueryExecutor,
+  {
+    project,
+    location,
+    prefixes,
+    minAgeHours = DEFAULT_STALE_DATASET_MIN_AGE_HOURS,
+  }: IDropStaleDatasetsOptions,
+): Promise<void> {
+  if (prefixes.length === 0) {
+    return;
+  }
+  for (const prefix of prefixes) {
+    // Guards against SQL injection, and against an empty prefix matching every test's datasets.
+    if (!DATASET_PREFIX_REGEX.test(prefix)) {
+      throw new Error(`Invalid dataset name prefix: "${prefix}"`);
+    }
+  }
+  const namePattern = `^(${prefixes.join("|")})_${UNIQUE_SUFFIX_PATTERN}$`;
+
+  let staleDatasets: string[];
+  try {
+    const { rows } = await executor.execute(
+      `select schema_name
+       from \`${project}\`.\`region-${location.toLowerCase()}\`.INFORMATION_SCHEMA.SCHEMATA
+       where regexp_contains(schema_name, r'${namePattern}')
+         and creation_time < timestamp_sub(current_timestamp(), interval ${minAgeHours} hour)`,
+    );
+    staleDatasets = rows.map((row: { schema_name: string }) => row.schema_name);
+  } catch (e) {
+    console.warn(`Failed to list stale datasets: ${e?.message || e}`);
+    return;
+  }
+  // Another run may be dropping the same datasets concurrently, which is fine.
+  await dropDatasets(executor, project, staleDatasets);
+}

testing/fixtures.ts

@@ -5,6 +5,8 @@ import { IHookHandler } from "df/testing";
 
 // TmpDirFixture can be used to create unique temporary directories which will be cleaned up
 // at the end of a test run. Intended for use within bazel tests.
+// In a { parallel: true } suite, pass `after`, not `afterEach`: the cleanup deletes the directories
+// of all tests, including ones that are still running.
 export class TmpDirFixture {
   private static dirCounter = 0;
 

tests/integration/bigquery.spec.ts

@@ -3,15 +3,46 @@ import { expect } from "chai";
 import * as dfapi from "df/cli/api";
 import { BigQueryDbAdapter } from "df/cli/api/dbadapters/bigquery";
 import { ExecutionSql } from "df/cli/api/dbadapters/execution_sql";
-import { CREDENTIALS_PATH } from "df/cli/index_test_base";
+import {
+  CREDENTIALS_PATH,
+  INTEGRATION_TEST_LOCATION,
+  INTEGRATION_TEST_PROJECT,
+} from "df/cli/index_test_base";
 import { dataform } from "df/protos/ts";
 import { suite, test } from "df/testing";
+import { dropDatasets, dropStaleDatasets, makeUniqueSuffix } from "df/testing/bigquery";
 import { integrationSuite } from "df/testing/credentials";
-import { cleanWarehouse, compile } from "df/tests/integration/utils";
+import { compile, datasetsOf } from "df/tests/integration/utils";
 
-integrationSuite("@dataform/integration/bigquery", { parallel: true }, () => {
+const SEARCH_SCHEMA_SUFFIX = "search";
+const SEARCH_DATASET = `df_integration_test_${SEARCH_SCHEMA_SUFFIX}`;
+
+integrationSuite("@dataform/integration/bigquery", { parallel: true }, ({ before, after }) => {
   const credentials = dfapi.credentials.read(CREDENTIALS_PATH);
   const dbadapter = new BigQueryDbAdapter(credentials);
+  const uniqueSuffix = makeUniqueSuffix();
+
+  // The datasets created by the tests, without the unique suffix.
+  let datasets: string[];
+
+  before("drop stale datasets", async () => {
+    datasets = datasetsOf(
+      await compile("tests/integration/bigquery_project", SEARCH_SCHEMA_SUFFIX),
+    );
+    await dropStaleDatasets(dbadapter, {
+      project: INTEGRATION_TEST_PROJECT,
+      location: INTEGRATION_TEST_LOCATION,
+      prefixes: datasets,
+    });
+  });
+
+  after("drop datasets", async () => {
+    await dropDatasets(
+      dbadapter,
+      INTEGRATION_TEST_PROJECT,
+      datasets.map((dataset) => `${dataset}_${uniqueSuffix}`),
+    );
+  });
 
   test("run unit tests", async () => {
     const compiledGraph = await compile("tests/integration/bigquery_project", "unit_tests");
@@ -86,10 +117,11 @@ integrationSuite("@dataform/integration/bigquery", { parallel: true }, () => {
   });
 
   test("search", async () => {
-    const compiledGraph = await compile("tests/integration/bigquery_project", "search");
-
-    // Drop all the tables before we do anything.
-    await cleanWarehouse(compiledGraph, dbadapter);
+    const compiledGraph = await compile(
+      "tests/integration/bigquery_project",
+      `${SEARCH_SCHEMA_SUFFIX}_${uniqueSuffix}`,
+    );
+    const dataset = `${SEARCH_DATASET}_${uniqueSuffix}`;
 
     // Run the project.
     const executionGraph = await dfapi.build(
@@ -106,8 +138,9 @@ integrationSuite("@dataform/integration/bigquery", { parallel: true }, () => {
     );
 
     const [fullSearch, partialSearch, columnSearch] = await Promise.all([
-      dbadapter.search("df_integration_test_search"),
-      dbadapter.search("test_sear"),
+      dbadapter.search(dataset),
+      // A partial match, that still only matches this run's dataset.
+      dbadapter.search(`test_search_${uniqueSuffix}`),
       dbadapter.search("val"),
     ]);
 

tests/integration/evaluate.spec.ts

@@ -2,20 +2,57 @@ import { expect } from "chai";
 
 import * as dfapi from "df/cli/api";
 import { BigQueryDbAdapter } from "df/cli/api/dbadapters/bigquery";
-import { CREDENTIALS_PATH, INTEGRATION_TEST_PROJECT } from "df/cli/index_test_base";
+import {
+  CREDENTIALS_PATH,
+  INTEGRATION_TEST_LOCATION,
+  INTEGRATION_TEST_PROJECT,
+} from "df/cli/index_test_base";
 import { targetAsReadableString } from "df/core/targets";
 import { dataform } from "df/protos/ts";
 import { test } from "df/testing";
+import { dropDatasets, dropStaleDatasets, makeUniqueSuffix } from "df/testing/bigquery";
 import { integrationSuite } from "df/testing/credentials";
-import { compile, keyBy } from "df/tests/integration/utils";
+import { compile, datasetsOf, keyBy } from "df/tests/integration/utils";
 
-integrationSuite("@dataform/integration/evaluate", () => {
+const EVALUATE_SCHEMA_SUFFIX = "evaluate";
+const EVALUATE_DATASET = `df_integration_test_${EVALUATE_SCHEMA_SUFFIX}`;
+const EVALUATE_ASSERTIONS_DATASET = `df_integration_test_assertions_${EVALUATE_SCHEMA_SUFFIX}`;
+
+integrationSuite("@dataform/integration/evaluate", ({ before, after }) => {
   const credentials = dfapi.credentials.read(CREDENTIALS_PATH);
   const dbadapter = new BigQueryDbAdapter(credentials);
+  const uniqueSuffix = makeUniqueSuffix();
+
+  // The datasets created by the tests, without the unique suffix.
+  let datasets: string[];
+
+  before("drop stale datasets", async () => {
+    datasets = datasetsOf(
+      await compile("tests/integration/bigquery_project", EVALUATE_SCHEMA_SUFFIX),
+    );
+    await dropStaleDatasets(dbadapter, {
+      project: INTEGRATION_TEST_PROJECT,
+      location: INTEGRATION_TEST_LOCATION,
+      prefixes: datasets,
+    });
+  });
+
+  after("drop datasets", async () => {
+    await dropDatasets(
+      dbadapter,
+      INTEGRATION_TEST_PROJECT,
+      datasets.map((dataset) => `${dataset}_${uniqueSuffix}`),
+    );
+  });
 
   test("evaluate from valid compiled graph as valid", async () => {
     // Create and run the project.
-    const compiledGraph = await compile("tests/integration/bigquery_project", "evaluate");
+    const compiledGraph = await compile(
+      "tests/integration/bigquery_project",
+      `${EVALUATE_SCHEMA_SUFFIX}_${uniqueSuffix}`,
+    );
+    const dataset = `${EVALUATE_DATASET}_${uniqueSuffix}`;
+    const assertionsDataset = `${EVALUATE_ASSERTIONS_DATASET}_${uniqueSuffix}`;
     const executionGraph = await dfapi.build(compiledGraph, {}, dbadapter);
     await dfapi.run(dbadapter, executionGraph).result();
 
@@ -27,44 +64,34 @@ integrationSuite("@dataform/integration/evaluate", () => {
       targetAsReadableString(t.target),
     );
 
-    const view =
-      tablesByName[`${INTEGRATION_TEST_PROJECT}.df_integration_test_evaluate.example_view`];
+    const view = tablesByName[`${INTEGRATION_TEST_PROJECT}.${dataset}.example_view`];
     let evaluations = await dbadapter.evaluate(dataform.Table.create(view));
     expect(evaluations.length).to.equal(1);
     expect(evaluations[0].status).to.equal(dataform.QueryEvaluation.QueryEvaluationStatus.SUCCESS);
 
     const materializedView =
-      tablesByName[
-        `${INTEGRATION_TEST_PROJECT}.df_integration_test_evaluate.example_materialized_view`
-      ];
+      tablesByName[`${INTEGRATION_TEST_PROJECT}.${dataset}.example_materialized_view`];
     evaluations = await dbadapter.evaluate(dataform.Table.create(materializedView));
     expect(evaluations.length).to.equal(1);
     expect(evaluations[0].status).to.equal(dataform.QueryEvaluation.QueryEvaluationStatus.SUCCESS);
 
-    const table =
-      tablesByName[`${INTEGRATION_TEST_PROJECT}.df_integration_test_evaluate.example_table`];
+    const table = tablesByName[`${INTEGRATION_TEST_PROJECT}.${dataset}.example_table`];
     evaluations = await dbadapter.evaluate(dataform.Table.create(table));
     expect(evaluations.length).to.equal(1);
     expect(evaluations[0].status).to.equal(dataform.QueryEvaluation.QueryEvaluationStatus.SUCCESS);
 
-    const operation =
-      operationsByName[
-        `${INTEGRATION_TEST_PROJECT}.df_integration_test_evaluate.example_operation`
-      ];
+    const operation = operationsByName[`${INTEGRATION_TEST_PROJECT}.${dataset}.example_operation`];
     evaluations = await dbadapter.evaluate(dataform.Operation.create(operation));
     expect(evaluations.length).to.equal(1);
     expect(evaluations[0].status).to.equal(dataform.QueryEvaluation.QueryEvaluationStatus.SUCCESS);
 
     const assertion =
-      assertionsByName[
-        `${INTEGRATION_TEST_PROJECT}.df_integration_test_assertions_evaluate.example_assertion_pass`
-      ];
+      assertionsByName[`${INTEGRATION_TEST_PROJECT}.${assertionsDataset}.example_assertion_pass`];
     evaluations = await dbadapter.evaluate(dataform.Assertion.create(assertion));
     expect(evaluations.length).to.equal(1);
     expect(evaluations[0].status).to.equal(dataform.QueryEvaluation.QueryEvaluationStatus.SUCCESS);
 
-    const incremental =
-      tablesByName[`${INTEGRATION_TEST_PROJECT}.df_integration_test_evaluate.example_incremental`];
+    const incremental = tablesByName[`${INTEGRATION_TEST_PROJECT}.${dataset}.example_incremental`];
     evaluations = await dbadapter.evaluate(dataform.Table.create(incremental));
     expect(evaluations.length).to.equal(2);
     expect(evaluations[0].status).to.equal(dataform.QueryEvaluation.QueryEvaluationStatus.SUCCESS);