Database function (#1438)

* Refactored name() function to finalize name * Added database() function * Bumped minor version to 2.3.0 * Refactored database() to add an error when the database is undefined * Inlined database & ref error messages
dataform-co · Jan 30, 2023 · 8f18e36 · 8f18e36
1 parent de78a65
commit 8f18e36
Show file tree

Hide file tree

Showing 7 changed files with 149 additions and 9 deletions.
diff --git a/core/assertion.ts b/core/assertion.ts
@@ -195,8 +195,7 @@ export class AssertionContext implements ICommonContext {
   public ref(ref: Resolvable | string[], ...rest: string[]) {
     ref = toResolvable(ref, rest);
     if (!resolvableAsTarget(ref)) {
-      const message = `Action name is not specified`;
-      this.assertion.session.compileError(new Error(message));
+      this.assertion.session.compileError(new Error(`Action name is not specified`));
       return "";
     }
     this.assertion.dependencies(ref);
@@ -213,6 +212,19 @@ export class AssertionContext implements ICommonContext {
     );
   }
 
+  public database(): string {
+    if (!this.assertion.proto.target.database) {
+      this.assertion.session.compileError(
+        new Error(`Warehouse does not support multiple databases`)
+      );
+      return "";
+    }
+
+    return this.assertion.session.finalizeDatabase(
+      this.assertion.proto.target.database
+    );
+  }
+
   public dependencies(name: Resolvable | Resolvable[]) {
     this.assertion.dependencies(name);
     return "";

diff --git a/core/common.ts b/core/common.ts
@@ -50,6 +50,11 @@ export interface ICommonContext {
    * Returns the schema of this dataset.
    */
   schema: () => string;
+
+  /**
+   * Returns the database of this dataset, if applicable.
+   */
+  database: () => string;
 }
 
 /**

diff --git a/core/operation.ts b/core/operation.ts
@@ -227,8 +227,7 @@ export class OperationContext implements ICommonContext {
   public ref(ref: Resolvable | string[], ...rest: string[]) {
     ref = toResolvable(ref, rest);
     if (!resolvableAsTarget(ref)) {
-      const message = `Action name is not specified`;
-      this.operation.session.compileError(new Error(message));
+      this.operation.session.compileError(new Error(`Action name is not specified`));
       return "";
     }
     this.operation.dependencies(ref);
@@ -245,6 +244,19 @@ export class OperationContext implements ICommonContext {
     );
   }
 
+  public database(): string {
+    if (!this.operation.proto.target.database) {
+      this.operation.session.compileError(
+        new Error(`Warehouse does not support multiple databases`)
+      );
+      return "";
+    }
+
+    return this.operation.session.finalizeDatabase(
+      this.operation.proto.target.database
+    );
+  }
+
   public dependencies(name: Resolvable | Resolvable[]) {
     this.operation.dependencies(name);
     return "";

diff --git a/core/table.ts b/core/table.ts
@@ -780,8 +780,7 @@ export class TableContext implements ITableContext {
   public ref(ref: Resolvable | string[], ...rest: string[]): string {
     ref = toResolvable(ref, rest);
     if (!resolvableAsTarget(ref)) {
-      const message = `Action name is not specified`;
-      this.table.session.compileError(new Error(message));
+      this.table.session.compileError(new Error(`Action name is not specified`));
       return "";
     }
     this.table.dependencies(ref);
@@ -793,7 +792,20 @@ export class TableContext implements ITableContext {
   }
 
   public schema(): string {
-    return this.table.session.finalizeSchema(this.table.proto.target.schema);
+    return this.table.session.finalizeSchema(
+      this.table.proto.target.schema
+    );
+  }
+
+  public database(): string {
+    if (!this.table.proto.target.database) {
+      this.table.session.compileError(new Error(`Warehouse does not support multiple databases`));
+      return "";
+    }
+
+    return this.table.session.finalizeDatabase(
+      this.table.proto.target.database
+    );
   }
 
   public type(type: TableType) {

diff --git a/core/test.ts b/core/test.ts
@@ -182,6 +182,10 @@ class RefReplacingContext implements ITableContext {
     return "";
   }
 
+  public database() {
+    return "";
+  }
+
   public where(where: Contextable<ITableContext, string>) {
     return "";
   }

diff --git a/tests/core/core.spec.ts b/tests/core/core.spec.ts
@@ -31,6 +31,16 @@ class TestConfigs {
     defaultLocation: "US"
   };
 
+  public static bigqueryWithDatabase: dataform.IProjectConfig = {
+    ...TestConfigs.bigquery,
+    defaultDatabase: "test-db",
+  };
+
+  public static bigqueryWithDatabaseAndSuffix: dataform.IProjectConfig = {
+    ...TestConfigs.bigqueryWithDatabase,
+    databaseSuffix: "suffix",
+  };
+
   public static snowflake: dataform.IProjectConfig = {
     warehouse: "snowflake",
     defaultSchema: "schema"
@@ -923,6 +933,38 @@ suite("@dataform/core", () => {
         expect(testTable.query).deep.equals(name)
       });
     });
+
+    [
+      {testConfig: TestConfigs.bigqueryWithDatabase, target: 'test-db.schema.test', database: 'test-db'},
+      {testConfig: TestConfigs.bigqueryWithDatabaseAndSuffix, target: 'test-db_suffix.schema.test', database: 'test-db_suffix'},
+    ].forEach(({testConfig, target, database}) => {
+      test(`database/suffix: "${target}"`, () => {
+        const session = new Session(path.dirname(__filename), testConfig);
+        session.publish("test", {type: "table"})
+          .query(ctx => ctx.database());
+
+        const graph = session.compile();
+
+        const testTable = graph.tables
+          .find(table => targetAsReadableString(table.target) === target);
+
+        expect(testTable.query).deep.equals(database)
+      });
+    });
+
+    test(`database fails when undefined`, () => {
+      const session = new Session(path.dirname(__filename), TestConfigs.redshift);
+      session.publish("test", {type: "table"}).query(ctx => ctx.database());
+
+      const graph = session.compile();
+
+      const testTable = graph.tables
+        .find(table => targetAsReadableString(table.target) === 'schema.test');
+
+      expect(graph.graphErrors.compilationErrors[0].message).deep
+        .equals("Warehouse does not support multiple databases");
+      expect(testTable.query).deep.equals("");
+    });
   });
 
   suite("resolve", () => {
@@ -994,7 +1036,33 @@ suite("@dataform/core", () => {
 
         expect(graph.operations[0].queries).deep.equals([finalizedName]);
       });
-    });    
+    });
+
+    [
+      {testConfig: TestConfigs.bigqueryWithDatabase, finalizedDatabase: 'test-db'},
+      {testConfig: TestConfigs.bigqueryWithDatabaseAndSuffix, finalizedDatabase: 'test-db_suffix'},
+    ].forEach(({testConfig, finalizedDatabase}) => {
+      test(`database with suffix: "${finalizedDatabase}"`, () => {
+        const session = new Session(path.dirname(__filename), testConfig);
+        session.operate("operate-1", ctx => ctx.database()).hasOutput(true);
+
+        const graph = session.compile();
+
+        expect(graph.operations[0].queries).deep.equals([finalizedDatabase]);
+      });
+    });
+
+    test(`database fails when undefined`, () => {
+      const session = new Session(path.dirname(__filename), TestConfigs.redshift);
+
+      session.operate("operate-1", ctx => ctx.database()).hasOutput(true);
+
+      const graph = session.compile();
+
+      expect(graph.graphErrors.compilationErrors[0].message).deep
+        .equals("Warehouse does not support multiple databases");
+      expect(JSON.stringify(graph.operations[0].queries)).deep.equals('[""]');
+    });
   });
 
   suite("graph", () => {
@@ -1269,5 +1337,32 @@ select '\${\`bar\`}'
         expect(JSON.stringify(graph.assertions[0].query)).to.deep.equal(`"${finalizedName}"`);
       });
     });
+
+    [
+      {testConfig: TestConfigs.bigqueryWithDatabase, finalizedDatabase: 'test-db'},
+      {testConfig: TestConfigs.bigqueryWithDatabaseAndSuffix, finalizedDatabase: 'test-db_suffix'},
+    ].forEach(({testConfig, finalizedDatabase}) => {
+      test(`database: ${finalizedDatabase}`, () => {
+        const session = new Session(path.dirname(__filename), {...testConfig, defaultDatabase: 'test-db'});
+
+        session.assert("database", ctx => ctx.database());
+
+        const graph = session.compile();
+
+        expect(JSON.stringify(graph.assertions[0].query)).to.deep.equal(`"${finalizedDatabase}"`);
+      });
+    });
+
+    test(`database fails when undefined`, () => {
+      const session = new Session(path.dirname(__filename), TestConfigs.redshift);
+
+      session.assert("database", ctx => ctx.database());
+
+      const graph = session.compile();
+
+      expect(graph.graphErrors.compilationErrors[0].message).deep
+        .equals("Warehouse does not support multiple databases");
+      expect(JSON.stringify(graph.assertions[0].query)).to.deep.equal('""');
+    });
   });
 });
diff --git a/version.bzl b/version.bzl
@@ -1,3 +1,3 @@
 # NOTE: If you change the format of this line, you must change the bash command
 # in /scripts/publish to extract the version string correctly.
-DF_VERSION = "2.2.0"
+DF_VERSION = "2.3.0"