Add Impersonate Service Account argument

cli/BUILD

@@ -41,7 +41,6 @@ ts_library(
         "@npm//@types/yargs",
         "@npm//chokidar",
         "@npm//glob",
-        "@npm//parse-duration",
         "@npm//readline-sync",
         "@npm//untildify",
         "@npm//yargs",

cli/api/BUILD

@@ -40,6 +40,7 @@ ts_library(
         "@npm//deepmerge",
         "@npm//fs-extra",
         "@npm//glob",
+        "@npm//google-auth-library",
         "@npm//google-sql-syntax-ts",
         "@npm//js-beautify",
         "@npm//js-yaml",

cli/api/dbadapters/bigquery.ts

@@ -1,4 +1,5 @@
-import { BigQuery, GetTablesResponse, TableField, TableMetadata } from "@google-cloud/bigquery";
+import { BigQuery, BigQueryOptions, GetTablesResponse, TableField, TableMetadata } from "@google-cloud/bigquery";
+import { GoogleAuth, Impersonated } from "google-auth-library";
 import Long from "long";
 import { PromisePoolExecutor } from "promise-pool-executor";
 
@@ -17,7 +18,9 @@ import { coerceAsError } from "df/common/errors/errors";
 import { retry } from "df/common/promises";
 import { dataform } from "df/protos/ts";
 
+const GOOGLE_CLOUD_PLATFORM_SCOPE = "https://www.googleapis.com/auth/cloud-platform";
 const EXTRA_GOOGLE_SCOPES = ["https://www.googleapis.com/auth/drive"];
+const IMPERSONATION_GOOGLE_SCOPES = [GOOGLE_CLOUD_PLATFORM_SCOPE, ...EXTRA_GOOGLE_SCOPES];
 
 const BIGQUERY_DATE_RELATED_FIELDS = [
   "BigQueryDate",
@@ -37,24 +40,42 @@ export interface IBigQueryExecutionOptions {
   reservation?: string;
 }
 
-export type BigQueryClientProvider = (projectId?: string) => BigQuery;
+export type BigQueryClientProvider = (projectId?: string) => BigQuery | Promise<BigQuery>;
 
 export function createBigQueryClientProvider(
   credentials: dataform.IBigQuery
 ): BigQueryClientProvider {
   const clients = new Map<string, BigQuery>();
-  return (projectId?: string) => {
+  return async (projectId?: string) => {
     projectId = projectId || credentials.projectId;
     if (!clients.has(projectId)) {
-      clients.set(
+      const clientConfig: BigQueryOptions = {
         projectId,
-        new BigQuery({
+        location: credentials.location
+      };
+
+      if (credentials.impersonateServiceAccount) {
+        // Impersonation requires cloud-platform for the source and target credentials, while the
+        // Drive scope remains necessary for BigQuery external tables backed by Google Drive.
+        const sourceAuth = new GoogleAuth({
           projectId,
-          scopes: EXTRA_GOOGLE_SCOPES,
-          location: credentials.location,
+          scopes: IMPERSONATION_GOOGLE_SCOPES,
           credentials: credentials.credentials && JSON.parse(credentials.credentials)
-        })
-      );
+        });
+
+        const authClient = await sourceAuth.getClient();
+
+        clientConfig.authClient = new Impersonated({
+          sourceClient: authClient,
+          targetPrincipal: credentials.impersonateServiceAccount,
+          targetScopes: IMPERSONATION_GOOGLE_SCOPES
+        });
+      } else {
+        clientConfig.scopes = EXTRA_GOOGLE_SCOPES;
+        clientConfig.credentials = credentials.credentials && JSON.parse(credentials.credentials);
+      }
+
+      clients.set(projectId, new BigQuery(clientConfig));
     }
     return clients.get(projectId);
   };
@@ -143,7 +164,7 @@ export class BigQueryDbAdapter implements IDbAdapter {
     return this.pool
       .addSingleTask({
         generator: async () => {
-          const [rows, , apiResponse] = await this.getClient().query({
+          const [rows, , apiResponse] = await (await this.getClient()).query({
             ...this.prepareQueryOptions(statement, options.rowLimit, options.bigquery, options.params),
             skipParsing: true
           } as any);
@@ -162,8 +183,8 @@ export class BigQueryDbAdapter implements IDbAdapter {
         try {
           await this.pool
             .addSingleTask({
-              generator: () =>
-                this.getClient().query({
+              generator: async () =>
+                (await this.getClient()).query({
                   useLegacySql: false,
                   query,
                   dryRun: true
@@ -189,11 +210,12 @@ export class BigQueryDbAdapter implements IDbAdapter {
 
   public async tables(database: string, schema?: string): Promise<dataform.ITableMetadata[]> {
     const datasetIds = schema ? [schema] : await this.schemas(database);
+    const client = await this.getClient(database);
     const tablesMetadata: dataform.ITableMetadata[] = [];
 
     await Promise.all(
       datasetIds.map(async datasetId => {
-        const [tables] = await this.getClient(database)
+        const [tables] = await client
           .dataset(datasetId)
           .getTables({ autoPaginate: true, maxResults: 1000 });
         await Promise.all(
@@ -286,14 +308,17 @@ export class BigQueryDbAdapter implements IDbAdapter {
   }
 
   public async deleteTable(target: dataform.ITarget): Promise<void> {
-    await this.getClient(target.database)
+    await (await this.getClient(target.database))
       .dataset(target.schema)
       .table(target.name)
       .delete({ ignoreNotFound: true });
   }
 
   public async schemas(database: string): Promise<string[]> {
-    const data = await this.getClient(database).getDatasets({ autoPaginate: true, maxResults: 1000 });
+    const data = await (await this.getClient(database)).getDatasets({
+      autoPaginate: true,
+      maxResults: 1000
+    });
     return data[0].map(dataset => dataset.id);
   }
 
@@ -313,7 +338,7 @@ export class BigQueryDbAdapter implements IDbAdapter {
       metadata.schema.fields
     );
 
-    await this.getClient(target.database)
+    await (await this.getClient(target.database))
       .dataset(target.schema)
       .table(target.name)
       .setMetadata({
@@ -325,7 +350,7 @@ export class BigQueryDbAdapter implements IDbAdapter {
 
   private async getMetadata(target: dataform.ITarget): Promise<TableMetadata> {
     try {
-      const table = await this.getClient(target.database)
+      const table = await (await this.getClient(target.database))
         .dataset(target.schema)
         .table(target.name)
         .getMetadata();
@@ -340,8 +365,8 @@ export class BigQueryDbAdapter implements IDbAdapter {
     }
   }
 
-  private getClient(projectId?: string) {
-    return this.clientProvider(projectId);
+  private async getClient(projectId?: string) {
+    return await this.clientProvider(projectId);
   }
 
   private async runQuery(
@@ -351,12 +376,13 @@ export class BigQueryDbAdapter implements IDbAdapter {
     byteLimit?: number,
     location?: string
   ): Promise<IExecutionResult> {
+    const client = await this.getClient();
     const results = await new Promise<any[]>((resolve, reject) => {
       const allRows = new LimitedResultSet({
         rowLimit,
         byteLimit
       });
-      const stream = this.getClient().createQueryStream({
+      const stream = client.createQueryStream({
         query,
         params,
         location
@@ -412,7 +438,8 @@ export class BigQueryDbAdapter implements IDbAdapter {
     return retry(
       async () => {
         try {
-          const job = await this.getClient().createQueryJob(
+          const client = await this.getClient();
+          const job = await client.createQueryJob(
             this.prepareQueryOptions(
               query,
               rowLimit,

cli/api/dbadapters/bigquery_test.ts

@@ -1,5 +1,6 @@
 import { Dataset, Table } from "@google-cloud/bigquery";
 import { expect } from "chai";
+import { Readable } from "stream";
 import { anything, instance, mock, verify, when } from "ts-mockito";
 
 import { BigQueryDbAdapter } from "df/cli/api/dbadapters/bigquery";
@@ -9,6 +10,8 @@ import { suite, test } from "df/testing";
 suite("BigQueryDbAdapter", () => {
   test("tables() with schema filters correctly", async () => {
     const mockBigQuery = mock<any>();
+    // ts-mockito invents a then() method, which prevents await from resolving this client.
+    when(mockBigQuery.then).thenReturn(undefined);
     const mockDataset = mock<Dataset>();
     const mockTable = mock<Table>();
 
@@ -49,6 +52,7 @@ suite("BigQueryDbAdapter", () => {
 
   test("tables() without schema lists all datasets and tables", async () => {
     const mockBigQuery = mock<any>();
+    when(mockBigQuery.then).thenReturn(undefined);
     const mockDataset = mock<Dataset>();
     const mockTable = mock<Table>();
     const schemaName = "schema1";
@@ -84,6 +88,41 @@ suite("BigQueryDbAdapter", () => {
     expect(result[0].target.name).to.equal(tableName);
   });
 
+  for (const asynchronous of [false, true]) {
+    test(`interactive queries accept ${asynchronous ? "asynchronous" : "synchronous"} clients`, async () => {
+      const client = {
+        createQueryStream: () => Readable.from([{ id: 1 }])
+      } as any;
+      const adapter = new BigQueryDbAdapter({ projectId: "project" }, {
+        clientProvider: () => asynchronous ? Promise.resolve(client) : client
+      });
+
+      const result = await adapter.execute("SELECT 1 AS id", { interactive: true });
+      expect(result.rows).deep.equals([{ id: 1 }]);
+    });
+  }
+
+  for (const failure of ["client initialization", "stream creation"]) {
+    test(`interactive queries reject on ${failure} failure`, async () => {
+      const expectedError = new Error(failure);
+      const adapter = new BigQueryDbAdapter({ projectId: "project" }, {
+        clientProvider: async () => {
+          if (failure === "client initialization") {
+            throw expectedError;
+          }
+          return {
+            createQueryStream: () => { throw expectedError; }
+          } as any;
+        }
+      });
+
+      await adapter.execute("SELECT 1", { interactive: true }).then(
+        () => { throw new Error("Expected the query to reject"); },
+        error => expect(error).equals(expectedError)
+      );
+    });
+  }
+
   test("setMetadata handles action without columns", async () => {
     // Partial mock for BigQuery client to avoid real network calls
     const mockBigQuery: any = {

cli/commands/run_command.ts

@@ -93,6 +93,9 @@ export const runCommand: ICommand<IRunArgs> = {
     const readCredentials = credentials.read(
       actuallyResolve(argv.projectDir, argv.credentials)
     );
+    if (argv.impersonateServiceAccount) {
+      readCredentials.impersonateServiceAccount = argv.impersonateServiceAccount;
+    }
 
     const dbadapter = new BigQueryDbAdapter(readCredentials);
     const executionGraph = await build(

cli/commands/run_options.ts

@@ -6,6 +6,8 @@ import {
   credentialsOption,
   IActionsArgs,
   ICredentialsArgs,
+  IImpersonateServiceAccountArgs,
+  impersonateServiceAccountOption,
   IJsonOutputArgs,
   IProjectDirArgs,
   ITimeoutArgs,
@@ -21,6 +23,7 @@ export interface IRunArgs
   extends IProjectDirArgs,
     IProjectConfigArgs,
     ICredentialsArgs,
+    IImpersonateServiceAccountArgs,
     IActionsArgs,
     IJsonOutputArgs,
     ITimeoutArgs {
@@ -173,6 +176,7 @@ export const runOptions: Array<INamedOption<yargs.Options, IRunArgs>> = [
   actionRetryLimitOption,
   actionsOption,
   credentialsOption,
+  impersonateServiceAccountOption,
   emitLineageOption,
   fullRefreshOption,
   includeDepsOption,

cli/commands/test_command.ts

@@ -5,6 +5,8 @@ import {
   assertProjectDirExists,
   credentialsOption,
   ICredentialsArgs,
+  IImpersonateServiceAccountArgs,
+  impersonateServiceAccountOption,
   IJsonOutputArgs,
   IProjectDirArgs,
   ITimeoutArgs,
@@ -26,6 +28,7 @@ import { ICommand } from "df/cli/yargswrapper";
 export interface ITestArgs
   extends IProjectDirArgs,
     ICredentialsArgs,
+    IImpersonateServiceAccountArgs,
     ITimeoutArgs,
     IJsonOutputArgs,
     IProjectConfigArgs {}
@@ -37,6 +40,7 @@ export const testCommand: ICommand<ITestArgs> = {
   check: [assertProjectDirExists],
   options: [
     credentialsOption,
+    impersonateServiceAccountOption,
     timeoutOption,
     jsonOutputOption,
     ...ProjectConfigOptions.allYargsOptions
@@ -60,6 +64,9 @@ export const testCommand: ICommand<ITestArgs> = {
     const readCredentials = credentials.read(
       actuallyResolve(argv.projectDir, argv.credentials)
     );
+    if (argv.impersonateServiceAccount) {
+      readCredentials.impersonateServiceAccount = argv.impersonateServiceAccount;
+    }
 
     if (!compiledGraph.tests.length) {
       printError("No unit tests found.");

cli/common_options.ts

@@ -1,10 +1,9 @@
 import * as fs from "fs";
-import parseDuration from "parse-duration";
 import * as path from "path";
 import yargs from "yargs";
 
 import { CREDENTIALS_FILENAME } from "df/cli/api/commands/credentials";
-import { actuallyResolve, assertPathExists } from "df/cli/util";
+import { actuallyResolve, assertPathExists, parseCliDuration } from "df/cli/util";
 import { INamedOption } from "df/cli/yargswrapper";
 
 export interface IProjectDirArgs {
@@ -63,6 +62,18 @@ export const credentialsOption: INamedOption<yargs.Options, ICredentialsArgs> =
     actuallyResolve(argv.projectDir, argv.credentials)
 };
 
+export interface IImpersonateServiceAccountArgs {
+  impersonateServiceAccount?: string;
+}
+
+export const impersonateServiceAccountOption: INamedOption<yargs.Options, IImpersonateServiceAccountArgs> = {
+  name: "impersonate-service-account",
+  option: {
+    describe: "Service account email to impersonate during authentication.",
+    type: "string"
+  }
+};
+
 export interface IJsonOutputArgs {
   json: boolean;
 }
@@ -77,7 +88,7 @@ export const jsonOutputOption: INamedOption<yargs.Options, IJsonOutputArgs> = {
 };
 
 export const coerceTimeout = (rawTimeoutString: string | null) =>
-  rawTimeoutString ? parseDuration(rawTimeoutString) : null;
+  rawTimeoutString ? parseCliDuration(rawTimeoutString) : null;
 
 export interface ITimeoutArgs {
   timeout: number | null;