Add Impersonate Service Account argument

cli/api/BUILD

@@ -33,6 +33,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

@@ -2,6 +2,7 @@ import Long from "long";
 import { PromisePoolExecutor } from "promise-pool-executor";
 
 import { BigQuery, GetTablesResponse, TableField, TableMetadata } from "@google-cloud/bigquery";
+import { GoogleAuth, Impersonated } from "google-auth-library";
 import { collectEvaluationQueries, QueryOrAction } from "df/cli/api/dbadapters/execution_sql";
 import { IBigQueryError, IDbAdapter, IDbClient, IExecutionResult, OnCancel } from "df/cli/api/dbadapters/index";
 import { parseBigqueryEvalError } from "df/cli/api/utils/error_parsing";
@@ -102,8 +103,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
@@ -128,7 +129,7 @@ export class BigQueryDbAdapter implements IDbAdapter {
   }
 
   public async tables(): Promise<dataform.ITarget[]> {
-    const datasets = await this.getClient().getDatasets({ autoPaginate: true, maxResults: 1000 });
+    const datasets = await (await this.getClient()).getDatasets({ autoPaginate: true, maxResults: 1000 });
     const tables = await Promise.all(
       datasets[0].map(dataset => dataset.getTables({ autoPaginate: true, maxResults: 1000 }))
     );
@@ -218,7 +219,7 @@ export class BigQueryDbAdapter implements IDbAdapter {
   }
 
   public async schemas(database: string): Promise<string[]> {
-    const data = await this.getClient(database).getDatasets();
+    const data = await (await this.getClient(database)).getDatasets();
     return data[0].map(dataset => dataset.id);
   }
 
@@ -238,7 +239,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({
@@ -250,7 +251,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();
@@ -265,19 +266,36 @@ export class BigQueryDbAdapter implements IDbAdapter {
     }
   }
 
-  private getClient(projectId?: string) {
+  private async getClient(projectId?: string) {
     projectId = projectId || this.bigQueryCredentials.projectId;
     if (!this.clients.has(projectId)) {
-      this.clients.set(
+      let clientConfig: any = {
         projectId,
-        new BigQuery({
-          projectId,
-          scopes: EXTRA_GOOGLE_SCOPES,
-          location: this.bigQueryCredentials.location,
-          credentials:
-            this.bigQueryCredentials.credentials && JSON.parse(this.bigQueryCredentials.credentials)
+        scopes: EXTRA_GOOGLE_SCOPES,
+        location: this.bigQueryCredentials.location
+      };
+
+      if (this.bigQueryCredentials.impersonateServiceAccount) {
+        // For impersonation, create an Impersonated credential directly
+        const sourceAuth = new GoogleAuth({
+          scopes: ['https://www.googleapis.com/auth/cloud-platform'],
+          projectId: projectId,
+          credentials: this.bigQueryCredentials.credentials && JSON.parse(this.bigQueryCredentials.credentials),
+        });
+
+        const authClient = await sourceAuth.getClient();
+
+        clientConfig.authClient = new Impersonated({
+          sourceClient: authClient,
+          targetPrincipal: this.bigQueryCredentials.impersonateServiceAccount,
+          targetScopes: ['https://www.googleapis.com/auth/cloud-platform']
         })
-      );
+    } else {
+        clientConfig.credentials =
+          this.bigQueryCredentials.credentials && JSON.parse(this.bigQueryCredentials.credentials);
+      }
+
+      this.clients.set(projectId, new BigQuery(clientConfig));
     }
     return this.clients.get(projectId);
   }
@@ -289,12 +307,12 @@ export class BigQueryDbAdapter implements IDbAdapter {
     byteLimit?: number,
     location?: string
   ) {
-    const results = await new Promise<any[]>((resolve, reject) => {
+    const results = await new Promise<any[]>(async (resolve, reject) => {
       const allRows = new LimitedResultSet({
         rowLimit,
         byteLimit
       });
-      const stream = this.getClient().createQueryStream({
+      const stream = (await this.getClient()).createQueryStream({
         query,
         params,
         location
@@ -330,7 +348,7 @@ export class BigQueryDbAdapter implements IDbAdapter {
     return retry(
       async () => {
         try {
-          const job = await this.getClient().createQueryJob({
+          const job = await (await this.getClient()).createQueryJob({
             useLegacySql: false,
             jobPrefix: "dataform-" + (jobPrefix ? `${jobPrefix}-` : ""),
             query,

cli/index.ts

@@ -166,6 +166,14 @@ const jobPrefixOption: INamedOption<yargs.Options> = {
   }
 };
 
+const impersonateServiceAccountOption: INamedOption<yargs.Options> = {
+  name: "impersonate-service-account",
+  option: {
+    describe: "Service account email to impersonate during authentication.",
+    type: "string"
+  }
+};
+
 const quietCompileOption: INamedOption<yargs.Options> = {
   name: "quiet",
   option: {
@@ -418,7 +426,7 @@ export function runCli() {
         format: `test [${projectDirMustExistOption.name}]`,
         description: "Run the dataform project's unit tests.",
         positionalOptions: [projectDirMustExistOption],
-        options: [credentialsOption, timeoutOption, ...ProjectConfigOptions.allYargsOptions],
+        options: [credentialsOption, impersonateServiceAccountOption, timeoutOption, ...ProjectConfigOptions.allYargsOptions],
         processFn: async argv => {
           print("Compiling...\n");
           const compiledGraph = await compile({
@@ -434,6 +442,10 @@ export function runCli() {
           const readCredentials = credentials.read(
             getCredentialsPath(argv[projectDirOption.name], argv[credentialsOption.name])
           );
+          if (argv[impersonateServiceAccountOption.name]) {
+            (readCredentials as any).impersonateServiceAccount =
+              argv[impersonateServiceAccountOption.name];
+          }
 
           if (!compiledGraph.tests.length) {
             printError("No unit tests found.");
@@ -478,6 +490,7 @@ export function runCli() {
           },
           actionsOption,
           credentialsOption,
+          impersonateServiceAccountOption,
           fullRefreshOption,
           includeDepsOption,
           includeDependentsOption,
@@ -513,6 +526,10 @@ export function runCli() {
           const readCredentials = credentials.read(
             getCredentialsPath(argv[projectDirOption.name], argv[credentialsOption.name])
           );
+          if (argv[impersonateServiceAccountOption.name]) {
+            (readCredentials as any).impersonateServiceAccount =
+              argv[impersonateServiceAccountOption.name];
+          }
 
           const dbadapter = new BigQueryDbAdapter(readCredentials);
           const executionGraph = await build(

package.json

@@ -5,7 +5,7 @@
     "@bazel/labs": "^0.42.3",
     "@bazel/rollup": "^3.0.0",
     "@bazel/typescript": "^3.0.0",
-    "@google-cloud/bigquery": "^5.6.0",
+    "@google-cloud/bigquery": "~7.1.1",
     "@google-cloud/storage": "^5.8.2",
     "@rollup/plugin-node-resolve": "^7.1.3",
     "@types/chai": "^4.1.7",
@@ -36,6 +36,7 @@
     "estraverse": "^5.1.0",
     "fs-extra": "^9.0.0",
     "glob": "^10.3.3",
+    "google-auth-library": "~8.9.0",
     "google-sql-syntax-ts": "^1.0.3",
     "js-beautify": "^1.10.2",
     "js-yaml": "^4.1.0",

packages/@dataform/cli/BUILD

@@ -33,6 +33,7 @@ externals = [
     "deepmerge",
     "fs-extra",
     "glob",
+    "google-auth-library",
     "google-sql-syntax-ts",
     "js-beautify",
     "js-yaml",

protos/profiles.proto

@@ -11,6 +11,8 @@ message BigQuery {
   string credentials = 3;
   // Options are listed here: https://cloud.google.com/bigquery/docs/locations
   string location = 4;
+  // Service account email to impersonate during authentication
+  string impersonate_service_account = 5;
 
   reserved 2;
 }