diff --git a/cli/api/BUILD b/cli/api/BUILD index 94d7bcdf2..675f2b052 100644 --- a/cli/api/BUILD +++ b/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", diff --git a/cli/api/dbadapters/bigquery.ts b/cli/api/dbadapters/bigquery.ts index fd2193f7f..b58e18d98 100644 --- a/cli/api/dbadapters/bigquery.ts +++ b/cli/api/dbadapters/bigquery.ts @@ -1,3 +1,4 @@ +import { GoogleAuth, Impersonated } from "google-auth-library"; import Long from "long"; import { PromisePoolExecutor } from "promise-pool-executor"; @@ -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,8 @@ export class BigQueryDbAdapter implements IDbAdapter { } public async tables(): Promise { - const datasets = await this.getClient().getDatasets({ autoPaginate: true, maxResults: 1000 }); + const client = await this.getClient(); + const datasets = await client.getDatasets({ autoPaginate: true, maxResults: 1000 }); const tables = await Promise.all( datasets[0].map(dataset => dataset.getTables({ autoPaginate: true, maxResults: 1000 })) ); @@ -218,7 +220,7 @@ export class BigQueryDbAdapter implements IDbAdapter { } public async schemas(database: string): Promise { - const data = await this.getClient(database).getDatasets(); + const data = await (await this.getClient(database)).getDatasets(); return data[0].map(dataset => dataset.id); } @@ -238,7 +240,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 +252,7 @@ export class BigQueryDbAdapter implements IDbAdapter { private async getMetadata(target: dataform.ITarget): Promise { 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 +267,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( + const clientConfig: any = { projectId, - new BigQuery({ + scopes: EXTRA_GOOGLE_SCOPES, + location: this.bigQueryCredentials.location + }; + + if (this.bigQueryCredentials.impersonateServiceAccount) { + // For impersonation, create an Impersonated credential directly + const sourceAuth = new GoogleAuth({ projectId, - scopes: EXTRA_GOOGLE_SCOPES, - location: this.bigQueryCredentials.location, - credentials: - this.bigQueryCredentials.credentials && JSON.parse(this.bigQueryCredentials.credentials) + scopes: ['https://www.googleapis.com/auth/cloud-platform'], + 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 +308,12 @@ export class BigQueryDbAdapter implements IDbAdapter { byteLimit?: number, location?: string ) { - const results = await new Promise((resolve, reject) => { + const results = await new Promise(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 +349,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({ useLegacySql: false, jobPrefix: "dataform-" + (jobPrefix ? `${jobPrefix}-` : ""), query, diff --git a/cli/index.ts b/cli/index.ts index 05c014af9..29908339b 100644 --- a/cli/index.ts +++ b/cli/index.ts @@ -171,6 +171,14 @@ const jobPrefixOption: INamedOption = { } }; +const impersonateServiceAccountOption: INamedOption = { + name: "impersonate-service-account", + option: { + describe: "Service account email to impersonate during authentication.", + type: "string" + } +}; + const bigqueryJobLabelsOption: INamedOption = { name: "job-labels", option: { @@ -191,6 +199,7 @@ const bigqueryJobLabelsOption: INamedOption = { } }; + const quietCompileOption: INamedOption = { name: "quiet", option: { @@ -463,7 +472,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({ @@ -479,6 +488,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."); @@ -523,6 +536,7 @@ export function runCli() { }, actionsOption, credentialsOption, + impersonateServiceAccountOption, fullRefreshOption, includeDepsOption, includeDependentsOption, @@ -559,6 +573,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( diff --git a/package.json b/package.json index 6b7ca2f4d..488c1acff 100644 --- a/package.json +++ b/package.json @@ -36,6 +36,7 @@ "estraverse": "^5.1.0", "fs-extra": "^9.0.0", "glob": "^10.5.0", + "google-auth-library": "~8.9.0", "google-sql-syntax-ts": "^1.0.3", "js-beautify": "^1.10.2", "js-yaml": "^4.1.1", diff --git a/packages/@dataform/cli/BUILD b/packages/@dataform/cli/BUILD index 7a1d7085c..07243da07 100644 --- a/packages/@dataform/cli/BUILD +++ b/packages/@dataform/cli/BUILD @@ -33,6 +33,7 @@ externals = [ "deepmerge", "fs-extra", "glob", + "google-auth-library", "google-sql-syntax-ts", "js-beautify", "js-yaml", diff --git a/protos/profiles.proto b/protos/profiles.proto index 09eabcdb8..37b982ac9 100644 --- a/protos/profiles.proto +++ b/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; } diff --git a/yarn.lock b/yarn.lock index 53185630e..2781d8be8 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2076,7 +2076,7 @@ google-auth-library@^7.0.0, google-auth-library@^7.0.2: jws "^4.0.0" lru-cache "^6.0.0" -google-auth-library@^8.0.2: +google-auth-library@^8.0.2, google-auth-library@~8.9.0: version "8.9.0" resolved "https://registry.yarnpkg.com/google-auth-library/-/google-auth-library-8.9.0.tgz#15a271eb2ec35d43b81deb72211bd61b1ef14dd0" integrity sha512-f7aQCJODJFmYWN6PeNKzgvy9LI2tYmXnzpNDHEjG5sDNPgGb2FXQyTBnXeSH+PAtpKESFD+LmHw3Ox3mN7e1Fg== @@ -3961,8 +3961,7 @@ string_decoder@~1.1.1: dependencies: safe-buffer "~5.1.0" -"strip-ansi-cjs@npm:strip-ansi@^6.0.1", strip-ansi@^6.0.1: - name strip-ansi-cjs +"strip-ansi-cjs@npm:strip-ansi@^6.0.1": version "6.0.1" resolved "https://registry.yarnpkg.com/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9" integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A== @@ -3983,6 +3982,13 @@ strip-ansi@^6.0.0: dependencies: ansi-regex "^5.0.0" +strip-ansi@^6.0.1: + version "6.0.1" + resolved "https://registry.yarnpkg.com/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9" + integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A== + dependencies: + ansi-regex "^5.0.1" + strip-ansi@^7.0.1: version "7.1.0" resolved "https://registry.yarnpkg.com/strip-ansi/-/strip-ansi-7.1.0.tgz#d5b6568ca689d8561370b0707685d22434faff45"