diff --git a/cli/api/lineage/emitter.ts b/cli/api/lineage/emitter.ts index 6d664a40e..74685b55a 100644 --- a/cli/api/lineage/emitter.ts +++ b/cli/api/lineage/emitter.ts @@ -128,7 +128,8 @@ export class LineageEmitter { action, actionResult, projectId, - location + location, + this.credentials.projectId ); const client = this.clientProvider(projectId, GLOBAL_LINEAGE_ENDPOINT); diff --git a/cli/api/lineage/payload_builder.ts b/cli/api/lineage/payload_builder.ts index 6b5119a20..0b3894a17 100644 --- a/cli/api/lineage/payload_builder.ts +++ b/cli/api/lineage/payload_builder.ts @@ -2,10 +2,26 @@ import { createHash, randomUUID } from "crypto"; import Long from "long"; import * as path from "path"; +import { version } from "df/core/version"; import { dataform } from "df/protos/ts"; const PRODUCER_URL = "https://github.com/dataform-co/dataform"; +const NOMINAL_TIME_FACET_SCHEMA = + "https://openlineage.io/spec/facets/1-0-1/NominalTimeRunFacet.json"; +const PARENT_RUN_FACET_SCHEMA = + "https://openlineage.io/spec/facets/1-2-0/ParentRunFacet.json#/$defs/ParentRunFacet"; +const EXTERNAL_QUERY_FACET_SCHEMA = + "https://openlineage.io/spec/facets/1-0-2/ExternalQueryRunFacet.json"; +const ERROR_MESSAGE_FACET_SCHEMA = + "https://openlineage.io/spec/facets/1-0-1/ErrorMessageRunFacet.json"; +const SQL_JOB_FACET_SCHEMA = "https://openlineage.io/spec/facets/1-1-0/SQLJobFacet.json"; +const GCP_LINEAGE_JOB_FACET_SCHEMA = + "https://openlineage.io/spec/facets/1-0-0/GcpLineageJobFacet.json#/$defs/GcpLineageJobFacet"; +const JOB_TYPE_FACET_SCHEMA = + "https://openlineage.io/spec/facets/2-0-4/JobTypeJobFacet.json#/$defs/JobTypeJobFacet"; +const RUN_EVENT_SCHEMA = "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/RunEvent"; + /** * Assembles OpenLineage RunEvent payloads for Dataform actions. * @@ -32,7 +48,8 @@ export class LineagePayloadBuilder { action: dataform.IExecutionAction, actionResult: dataform.IActionResult, projectId: string, - location: string + location: string, + credentialsProjectId?: string ): { [key: string]: any } { if (!this.workdirHash && this.projectDir) { this.workdirHash = createHash("sha256").update(this.projectDir).digest("hex").slice(0, 16); @@ -78,7 +95,7 @@ export class LineagePayloadBuilder { const parentJobName = `${projectId}.${location}.cli.${workdirIdentifier}.run`; const nominalTime: any = { - _schemaURL: "https://openlineage.io/spec/facets/1-0-1/NominalTimeRunFacet.json", + _schemaURL: NOMINAL_TIME_FACET_SCHEMA, nominalStartTime: new Date( toMillis(actionResult.timing?.startTimeMillis) ?? Date.now() ).toISOString() @@ -93,8 +110,7 @@ export class LineagePayloadBuilder { nominalTime, parent: { _producer: PRODUCER_URL, - _schemaURL: - "https://openlineage.io/spec/facets/1-0-1/ParentRunFacet.json#/$defs/ParentRunFacet", + _schemaURL: PARENT_RUN_FACET_SCHEMA, job: { namespace: "dataform", name: parentJobName @@ -108,11 +124,48 @@ export class LineagePayloadBuilder { } }; + if (eventType !== "START") { + const bqJobId = extractBqJobId(actionResult); + if (bqJobId && credentialsProjectId) { + runFacets.externalQuery = { + _producer: PRODUCER_URL, + _schemaURL: EXTERNAL_QUERY_FACET_SCHEMA, + externalQueryId: `${credentialsProjectId}.${location}.${bqJobId}`, + source: "bigquery" + }; + } + } + + if (eventType === "FAIL") { + const errorMessages = actionResult.tasks + ?.map(t => t.errorMessage) + .filter(msg => !!msg) + .join("; "); + if (errorMessages) { + runFacets.errorMessage = { + _schemaURL: ERROR_MESSAGE_FACET_SCHEMA, + message: errorMessages, + programmingLanguage: "typescript" + }; + } + } + + const sqlStatements = action.tasks + ?.map(task => task.statement) + .filter(stmt => !!stmt) + .join(";\n"); + const jobFacets: any = {}; + if (sqlStatements) { + jobFacets.sql = { + _schemaURL: SQL_JOB_FACET_SCHEMA, + query: sqlStatements + }; + } + jobFacets.gcp_lineage = { _producer: PRODUCER_URL, - _schemaURL: - "https://openlineage.io/spec/facets/1-0-0/GcpLineageJobFacet.json#/$defs/GcpLineageJobFacet", + _schemaURL: GCP_LINEAGE_JOB_FACET_SCHEMA, displayName: `BigQuery Pipelines action ${canonicalActionTarget}`, origin: { name: `projects/${projectId}/locations/${location}/cli/${workdirIdentifier}`, @@ -120,6 +173,20 @@ export class LineagePayloadBuilder { } }; + jobFacets.jobType = { + _producer: PRODUCER_URL, + _schemaURL: JOB_TYPE_FACET_SCHEMA, + integration: "BIGQUERY_PIPELINES", + jobType: "ACTION", + processingType: "BATCH" + }; + + jobFacets.gcp_bq_pipelines_job = { + dataformCoreVersion: version, + actionType: action.type, + actionName: canonicalActionTarget + }; + return { eventType, eventTime, @@ -135,7 +202,7 @@ export class LineagePayloadBuilder { inputs, outputs, producer: PRODUCER_URL, - schemaURL: "https://openlineage.io/spec/1-0-2/OpenLineage.json#/definitions/RunEvent" + schemaURL: RUN_EVENT_SCHEMA }; } @@ -164,6 +231,19 @@ function toMillis(val: Long | number | undefined | null): number | undefined { return undefined; } +function extractBqJobId(actionResult: dataform.IActionResult): string | undefined { + if (!actionResult.tasks) { + return undefined; + } + for (let i = actionResult.tasks.length - 1; i >= 0; i--) { + const jobId = actionResult.tasks[i]?.metadata?.bigquery?.jobId; + if (jobId) { + return jobId; + } + } + return undefined; +} + export function toProtoStruct(obj: { [key: string]: any }): any { const fields: { [key: string]: any } = {}; for (const key of Object.keys(obj)) { diff --git a/cli/api/lineage/payload_builder_test.ts b/cli/api/lineage/payload_builder_test.ts index edccea13b..15715bb48 100644 --- a/cli/api/lineage/payload_builder_test.ts +++ b/cli/api/lineage/payload_builder_test.ts @@ -89,6 +89,92 @@ suite("LineagePayloadBuilder", () => { expect(originName).to.equal("projects/proj/locations/us/cli/unknown-workdir"); }); + test("externalQuery run facet appears only on terminal events with a bqJobId", () => { + const builder = new LineagePayloadBuilder("/tmp/proj"); + const start = builder.build(ACTION, START_RESULT, "proj", "us", "cred-proj"); + const complete = builder.build(ACTION, COMPLETE_RESULT, "proj", "us", "cred-proj"); + expect(start.run.facets.externalQuery).to.equal(undefined); + expect(complete.run.facets.externalQuery).to.deep.include({ + externalQueryId: "cred-proj.us.job-abc", + source: "bigquery" + }); + }); + + test("externalQuery omitted when credentialsProjectId not passed even with bqJobId", () => { + const builder = new LineagePayloadBuilder("/tmp/proj"); + const complete = builder.build(ACTION, COMPLETE_RESULT, "proj", "us"); + expect(complete.run.facets.externalQuery).to.equal(undefined); + }); + + test("errorMessage run facet appears only on FAIL with non-empty task errorMessage", () => { + const builder = new LineagePayloadBuilder("/tmp/proj"); + const failedWithMsg = builder.build( + ACTION, + dataform.ActionResult.create({ + status: dataform.ActionResult.ExecutionStatus.FAILED, + tasks: [{ errorMessage: "boom" }, { errorMessage: "kaboom" }] + }), + "proj", + "us" + ); + expect(failedWithMsg.run.facets.errorMessage).to.deep.include({ + message: "boom; kaboom", + programmingLanguage: "typescript" + }); + const failedWithoutMsg = builder.build( + ACTION, + dataform.ActionResult.create({ + status: dataform.ActionResult.ExecutionStatus.FAILED + }), + "proj", + "us" + ); + expect(failedWithoutMsg.run.facets.errorMessage).to.equal(undefined); + }); + + test("sql job facet emits joined task statements", () => { + const multiTaskAction = dataform.ExecutionAction.create({ + target: { database: "proj", schema: "schema", name: "table" }, + type: "table", + tasks: [{ statement: "CREATE TABLE t AS SELECT 1" }, { statement: "SELECT 2" }] + }); + const builder = new LineagePayloadBuilder("/tmp/proj"); + const payload = builder.build(multiTaskAction, START_RESULT, "proj", "us"); + expect(payload.job.facets.sql).to.deep.include({ + query: "CREATE TABLE t AS SELECT 1;\nSELECT 2" + }); + }); + + test("sql job facet omitted when action has no statements", () => { + const emptyAction = dataform.ExecutionAction.create({ + target: { database: "proj", schema: "schema", name: "table" }, + type: "table" + }); + const builder = new LineagePayloadBuilder("/tmp/proj"); + const payload = builder.build(emptyAction, START_RESULT, "proj", "us"); + expect(payload.job.facets.sql).to.equal(undefined); + }); + + test("jobType job facet declares BigQuery Pipelines integration", () => { + const builder = new LineagePayloadBuilder("/tmp/proj"); + const payload = builder.build(ACTION, START_RESULT, "proj", "us"); + expect(payload.job.facets.jobType).to.deep.include({ + integration: "BIGQUERY_PIPELINES", + jobType: "ACTION", + processingType: "BATCH" + }); + }); + + test("gcp_bq_pipelines_job facet carries dataformCoreVersion + action metadata", () => { + const builder = new LineagePayloadBuilder("/tmp/proj"); + const payload = builder.build(ACTION, START_RESULT, "proj", "us"); + expect(payload.job.facets.gcp_bq_pipelines_job).to.deep.include({ + actionType: "table", + actionName: "schema.table" + }); + expect(payload.job.facets.gcp_bq_pipelines_job.dataformCoreVersion).to.be.a("string"); + }); + test("inputs list mirrors action.dependencyTargets in bigquery namespace", () => { const actionWithDeps = dataform.ExecutionAction.create({ target: { database: "proj", schema: "schema", name: "table" },