@malloy-publisher/server 0.0.232 → 0.0.234

package/src/service/materialization_service.spec.ts

@@ -39,6 +39,23 @@ import {
 
 type MockRepo = sinon.SinonStubbedInstance<ResourceRepository>;
 
+/**
+ * Turn per-query metadata on for one test and put the environment back after.
+ * The feature ships dark, so a test that asserts what a statement carries has
+ * to enable it — and saying so per test keeps the default honest here: a test
+ * that forgets is testing the shipped default, which is "nothing attached".
+ */
+function withQueryMetadataOn(): void {
+   const original = process.env.PUBLISHER_QUERY_METADATA;
+   beforeEach(() => {
+      process.env.PUBLISHER_QUERY_METADATA = "on";
+   });
+   afterEach(() => {
+      if (original === undefined) delete process.env.PUBLISHER_QUERY_METADATA;
+      else process.env.PUBLISHER_QUERY_METADATA = original;
+   });
+}
+
 describe("redactConnectionSecrets", () => {
    it("strips credential values but keeps the message legible", () => {
       const source = {
@@ -877,6 +894,7 @@ describe("MaterializationService", () => {
 describe("deleteMaterialization telemetry", () => {
    let ctx: ReturnType<typeof createMocks>;
    let harness: MetricsHarness;
+   withQueryMetadataOn();
 
    beforeEach(async () => {
       ctx = createMocks();
@@ -936,6 +954,27 @@ describe("deleteMaterialization telemetry", () => {
       ).toBe(0);
    });
 
+   it("tags the retiring drops as ops work", async () => {
+      // A drop is warehouse work someone accounts for later, and it is the one
+      // statement an operator goes looking for after a table disappears.
+      const runSQL = sinon.stub().resolves();
+      dropTablesFixture(runSQL);
+
+      await ctx.service.deleteMaterialization("my-env", "pkg", "mat-1", {
+         dropTables: true,
+      });
+
+      expect(runSQL.callCount).toBeGreaterThan(0);
+      for (const call of runSQL.getCalls()) {
+         expect(call.args[1]?.queryMetadata).toMatchObject({
+            class: "ops",
+            environment: "my-env",
+            package: "pkg",
+            run_id: "mat-1",
+         });
+      }
+   });
+
    it("meters a failed drop as outcome=failure and still deletes the record", async () => {
       dropTablesFixture(sinon.stub().rejects(new Error("drop boom")));
 
@@ -1472,6 +1511,7 @@ describe("executeInstructedBuild", () => {
 
 describe("buildOneSource", () => {
    let ctx: ReturnType<typeof createMocks>;
+   withQueryMetadataOn();
    beforeEach(() => {
       ctx = createMocks();
    });
@@ -1481,12 +1521,19 @@ describe("buildOneSource", () => {
       physicalTableName: string,
       dialectName?: string,
       manifest: Manifest = new Manifest(),
+      metadata?: {
+         source?: Record<string, string>;
+         package?: Record<string, string>;
+         connection?: Record<string, string>;
+         context?: Record<string, string | undefined>;
+      },
    ): Promise<{ sourceEntityId: string; physicalTableName: string }> {
       const source = fakeSource({
          name: "orders",
          sourceEntityId: "abcdef1234567890",
          sql: "SELECT * FROM t",
          dialectName,
+         queryMetadata: metadata?.source,
       });
       const instruction: BuildInstruction = {
          sourceEntityId: "abcdef1234567890",
@@ -1504,6 +1551,7 @@ describe("buildOneSource", () => {
                m: Manifest,
                env: unknown,
                built: Record<string, unknown>,
+               buildMetadata?: unknown,
             ) => Promise<{ sourceEntityId: string; physicalTableName: string }>;
          }
       ).buildOneSource(
@@ -1514,14 +1562,85 @@ describe("buildOneSource", () => {
          manifest,
          {
             getApiConnection: () => {
+               if (metadata?.connection) {
+                  return { queryMetadata: metadata.connection };
+               }
                throw new Error("no connection config in this colocated test");
             },
             getEnvironmentPath: () => "/tmp/env",
          },
          {}, // builtEntries — colocated build with no prior storage upstreams
+         metadata
+            ? {
+                 packageMaterialization: metadata.package
+                    ? {
+                         schedule: null,
+                         freshness: null,
+                         queryMetadata: metadata.package,
+                      }
+                    : null,
+                 context: metadata.context ?? { queryClass: "materialize" },
+              }
+            : undefined,
       );
    }
 
+   it("tags every statement of the build with the resolved metadata", async () => {
+      // One unit of work: the warehouse's query history has to show the staging
+      // drop, the CTAS, the swap and the rename as the same build.
+      const runSQL = sinon.stub().resolves();
+      await callBuildOneSource({ runSQL }, "orders_v1", undefined, undefined, {
+         connection: { deployment: "eu" },
+         package: { team: "finance", tier: "bronze" },
+         source: { tier: "gold" },
+         context: { queryClass: "materialize", runId: "run-7" },
+      });
+
+      const bags = runSQL.getCalls().map((c) => c.args[1]?.queryMetadata);
+      expect(bags).toHaveLength(4);
+      for (const bag of bags) {
+         expect(bag).toEqual({
+            deployment: "eu",
+            team: "finance",
+            // The source's own declaration wins over the package's.
+            tier: "gold",
+            class: "materialize",
+            run_id: "run-7",
+            source: "orders",
+         });
+      }
+   });
+
+   it("issues untagged statements when nothing declares metadata", async () => {
+      // No metadata context (an internal caller) means the statements are exactly
+      // what they were before this existed.
+      const runSQL = sinon.stub().resolves();
+      await callBuildOneSource({ runSQL }, "orders_v1");
+      for (const call of runSQL.getCalls()) {
+         expect(call.args[1]?.queryMetadata).toBeUndefined();
+      }
+   });
+
+   it("tags the cleanup drop after a failed build", async () => {
+      // The cleanup statement belongs to the same build, and it is the one an
+      // operator goes looking for after a failure.
+      const runSQL = sinon.stub();
+      runSQL.onCall(0).resolves();
+      runSQL.onCall(1).rejects(new Error("create boom"));
+      runSQL.onCall(2).resolves();
+      await expect(
+         callBuildOneSource({ runSQL }, "orders_v1", undefined, undefined, {
+            source: { team: "finance" },
+            context: { queryClass: "materialize" },
+         }),
+      ).rejects.toThrow("create boom");
+      expect(runSQL.lastCall.args[1]?.queryMetadata).toEqual({
+         team: "finance",
+         class: "materialize",
+         source: "orders",
+      });
+   });
+
    it("stages, swaps, and renames in a crash-safe order", async () => {
       const runSQL = sinon.stub().resolves();
       const entry = await callBuildOneSource({ runSQL }, "orders_v1");

package/src/service/materialization_service.ts

@@ -43,7 +43,14 @@ import {
    deriveAnnotationFields,
    projectToPublicColumns,
    iterGraphSources,
+   resolveQueryMetadata,
 } from "./build_plan";
+import {
+   mergeQueryMetadata,
+   type QueryContext,
+   type QueryMetadata,
+} from "./query_metadata";
+import type { components } from "../api";
 import { getPersistStorageMode } from "../config";
 import { EnvironmentStore } from "./environment_store";
 import { assertMaterializationEligible } from "./materialization_eligibility";
@@ -88,6 +95,47 @@ interface BuildEnvironment {
    getEnvironmentPath(): string;
 }
 
+/**
+ * What a build needs to tag the statements it issues: the package-level layer of
+ * per-query metadata, and the run's own context. Assembled once per run — every
+ * field is constant for the run, so a source only adds its own name.
+ */
+interface BuildQueryMetadata {
+   packageMaterialization:
+      | components["schemas"]["PackageMaterializationConfig"]
+      | null;
+   context: QueryContext;
+}
+
+/**
+ * A connection's two metadata layers: the overridable default and the properties
+ * the deployment enforces.
+ *
+ * Fails open — a build must not fail because a connection's config could not be
+ * read for its tags — so an unreadable connection costs the layers, never the
+ * statement.
+ */
+function connectionMetadataLayers(
+   environment: BuildEnvironment,
+   connectionName: string,
+): { default: QueryMetadata | null; enforced: QueryMetadata | null } {
+   try {
+      const connection = environment.getApiConnection(connectionName);
+      return {
+         default: connection?.queryMetadata ?? null,
+         enforced: connection?.queryMetadataEnforced ?? null,
+      };
+   } catch (error) {
+      // Diagnosable rather than silent: the layer this costs is the enforced
+      // one, and it is the only drop with no metric behind it.
+      logger.debug("No query-metadata layers for connection", {
+         connectionName,
+         error,
+      });
+      return { default: null, enforced: null };
+   }
+}
+
 /**
  * Length of the sourceEntityId prefix used when synthesizing staging table
  * names. 12 hex chars is 48 bits of entropy, well inside every dialect's
@@ -413,6 +461,14 @@ export class MaterializationService {
           * metadata so a scheduled rebuild is distinguishable from a manual one.
           */
          trigger?: "ON_DEMAND" | "SCHEDULER";
+         /**
+          * What the caller knows about this run and the publisher does not,
+          * attached as query metadata to the statements the build issues. Its
+          * `trigger` also covers the case the publisher's own `trigger` cannot
+          * express (a publish), and its `runId` lets a caller's own id group the
+          * build's statements instead of the publisher's materialization id.
+          */
+         runContext?: components["schemas"]["RunContext"] | null;
       } = {},
    ): Promise<Materialization> {
       const environmentId = await this.resolveEnvironmentId(environmentName);
@@ -477,6 +533,7 @@ export class MaterializationService {
                referenceManifest: options.referenceManifest,
                strictUpstreams: options.strictUpstreams,
                trigger,
+               runContext: options.runContext ?? undefined,
             },
             signal,
          ),
@@ -504,6 +561,7 @@ export class MaterializationService {
          referenceManifest: ManifestReference[] | undefined;
          strictUpstreams: boolean | undefined;
          trigger: "ON_DEMAND" | "SCHEDULER";
+         runContext?: components["schemas"]["RunContext"];
       },
       signal: AbortSignal,
    ): Promise<void> {
@@ -620,6 +678,25 @@ export class MaterializationService {
             // Failure-path reclaim is ORCHESTRATED-ONLY on purpose — see
             // reclaimStorageTablesFromFailedRun.
             orchestrated ? { environmentId, packageName } : undefined,
+            {
+               // Optional for the same reason build_plan reads it optionally:
+               // callers that build from a lighter package surface still resolve,
+               // just without a package-level layer.
+               packageMaterialization: pkg.getMaterializationConfig?.() ?? null,
+               context: {
+                  queryClass: "materialize",
+                  environment: environmentName,
+                  package: packageName,
+                  // The caller's trigger wins because it can express a publish,
+                  // which the publisher's own two-value trigger cannot.
+                  trigger:
+                     opts.runContext?.trigger ?? opts.trigger?.toLowerCase(),
+                  // Default to the materialization id: the publisher always has a
+                  // run id, so a build's statements group in the backend's query
+                  // history whether or not the caller supplied one.
+                  runId: opts.runContext?.runId ?? id,
+               },
+            },
          );
 
          const sourcesBuilt = instructions.length;
@@ -1083,6 +1160,7 @@ export class MaterializationService {
       // Identity of the run, used only to reclaim storage tables this run created
       // if it fails part-way (see reclaimStorageTablesFromFailedRun).
       owner?: { environmentId: string; packageName: string },
+      buildMetadata?: BuildQueryMetadata,
    ): Promise<Record<string, ManifestEntry>> {
       const { graphs, sources, connectionDigests, connections } = compiled;
 
@@ -1200,6 +1278,7 @@ export class MaterializationService {
                   manifest,
                   environment,
                   entries,
+                  buildMetadata,
                );
                entries[sourceEntityId] = entry;
                if (entry.storageConnectionName) builtThisRun.push(entry);
@@ -1327,6 +1406,76 @@ export class MaterializationService {
       }
    }
 
+   /**
+    * The `RunSQLOptions` for one source's build statements: its resolved
+    * per-query metadata, merged under this run's context.
+    *
+    * Layers, least specific first: the executing connection's default, then what
+    * the model side declared for this source (package → model-file → `#@ persist`,
+    * already resolved by {@link resolveQueryMetadata}), then the run's context,
+    * which names the source. There is no request layer — a build has no request.
+    *
+    * Fails OPEN: a build must not fail because metadata could not be assembled,
+    * so an unresolvable connection just contributes no default, and a dropped
+    * property is logged and metered rather than thrown.
+    */
+   private buildRunSQLOptions(
+      persistSource: PersistSource,
+      environment: BuildEnvironment,
+      buildMetadata: BuildQueryMetadata | undefined,
+   ): { queryMetadata?: QueryMetadata } {
+      if (!buildMetadata) return {};
+      const connectionLayers = connectionMetadataLayers(
+         environment,
+         persistSource.connectionName,
+      );
+      const resolved = mergeQueryMetadata({
+         connection: connectionLayers.default,
+         enforced: connectionLayers.enforced,
+         model: resolveQueryMetadata(
+            persistSource,
+            buildMetadata.packageMaterialization,
+         ),
+         context: { ...buildMetadata.context, source: persistSource.name },
+      });
+      if (resolved.drops.length > 0) {
+         logger.warn("Dropped query-metadata properties for a build", {
+            sourceName: persistSource.name,
+            drops: resolved.drops,
+         });
+      }
+      return resolved.metadata ? { queryMetadata: resolved.metadata } : {};
+   }
+
+   /**
+    * The `RunSQLOptions` for the drops that retire a materialization's tables.
+    * No model layer: the source's declaration described how to BUILD it, and the
+    * source may no longer exist by the time its table is retired.
+    */
+   private dropRunSQLOptions(
+      environment: BuildEnvironment,
+      connectionName: string,
+      environmentName: string,
+      packageName: string,
+      materializationId: string,
+   ): { queryMetadata?: QueryMetadata } {
+      const connectionLayers = connectionMetadataLayers(
+         environment,
+         connectionName,
+      );
+      const resolved = mergeQueryMetadata({
+         connection: connectionLayers.default,
+         enforced: connectionLayers.enforced,
+         context: {
+            queryClass: "ops",
+            environment: environmentName,
+            package: packageName,
+            runId: materializationId,
+         },
+      });
+      return resolved.metadata ? { queryMetadata: resolved.metadata } : {};
+   }
+
    /**
     * Build a single instructed source into its assigned physical table.
     * COPY uses a staging table + atomic rename for crash-safety; the staging
@@ -1340,6 +1489,7 @@ export class MaterializationService {
       manifest: Manifest,
       environment: BuildEnvironment,
       builtEntries: Record<string, ManifestEntry>,
+      buildMetadata?: BuildQueryMetadata,
    ): Promise<ManifestEntry> {
       const sourceEntityId = instruction.sourceEntityId;
       const physicalTableName = instruction.physicalTableName;
@@ -1404,6 +1554,15 @@ export class MaterializationService {
          );
       }
 
+      // Every statement of this source's build carries the same metadata, so the
+      // warehouse's query history shows the staging CTAS, the drop and the rename
+      // as one attributable unit of work.
+      const runOptions = this.buildRunSQLOptions(
+         persistSource,
+         environment,
+         buildMetadata,
+      );
+
       const bareName = bareTableName(physicalTableName);
       const stagingTableName = `${physicalTableName}${stagingSuffix(sourceEntityId)}`;
       // The control plane sends the logical (unquoted) physical name; dialect-
@@ -1416,18 +1575,29 @@ export class MaterializationService {
       const quotedBareName = quoteIdentifier(bareName, dialect);
 
       const startTime = performance.now();
-      await connection.runSQL(`DROP TABLE IF EXISTS ${quotedStaging}`);
+      await connection.runSQL(
+         `DROP TABLE IF EXISTS ${quotedStaging}`,
+         runOptions,
+      );
       try {
          await connection.runSQL(
             `CREATE TABLE ${quotedStaging} AS (${buildSQL})`,
+            runOptions,
+         );
+         await connection.runSQL(
+            `DROP TABLE IF EXISTS ${quotedPhysical}`,
+            runOptions,
          );
-         await connection.runSQL(`DROP TABLE IF EXISTS ${quotedPhysical}`);
          await connection.runSQL(
             `ALTER TABLE ${quotedStaging} RENAME TO ${quotedBareName}`,
+            runOptions,
          );
       } catch (err) {
          try {
-            await connection.runSQL(`DROP TABLE IF EXISTS ${quotedStaging}`);
+            await connection.runSQL(
+               `DROP TABLE IF EXISTS ${quotedStaging}`,
+               runOptions,
+            );
          } catch (cleanupErr) {
             logger.warn(
                "Failed to clean up staging table after a failed build; physical leak",
@@ -2045,11 +2215,23 @@ export class MaterializationService {
             // successfully also drops successfully (container paths, hyphenated
             // BigQuery project ids, etc.).
             const dialect = connection.dialectName;
+            // Dropping a materialized table is warehouse work someone will have
+            // to account for later, so it is tagged like the build that created
+            // it — `ops` rather than `materialize`, because this is the
+            // lifecycle operation, not a build.
+            const dropOptions = this.dropRunSQLOptions(
+               environment,
+               connectionName,
+               environmentName,
+               packageName,
+               m.id,
+            );
             await connection.runSQL(
                `DROP TABLE IF EXISTS ${quoteTablePath(
                   physicalTableName,
                   dialect,
                )}`,
+               dropOptions,
             );
             // A crash between staging-create and rename can leave the staging
             // table behind; clean it up too while we hold the connection.
@@ -2058,6 +2240,7 @@ export class MaterializationService {
                   `${physicalTableName}${stagingSuffix(entry.sourceEntityId)}`,
                   dialect,
                )}`,
+               dropOptions,
             );
             recordDropTables("success", "in_warehouse");
             logger.info("Dropped materialized table on delete", {

package/src/service/materialization_test_fixtures.ts

@@ -80,34 +80,78 @@ interface FakeFreshnessSchedule {
    freshness?: { window?: string; fallback?: string };
 }
 
+/** A nested tag tree: scalar leaves and property collections. */
+type FakeTagTree = { [key: string]: string | FakeTagTree };
+
 /**
- * Build a fake Malloy `Tag` supporting both readers the build plan uses:
- * `entries()` (scalar `#@ persist` key=value pairs, for deriveAnnotationFields)
- * and the path-based `text(...at)` (dotted `freshness.window`, for
- * resolveFreshness).
+ * Build a fake Malloy `Tag` over a nested tree, supporting all three readers the
+ * build plan uses: `entries()` (scalar `#@ persist` key=value pairs, for
+ * deriveAnnotationFields), the path-based `text(...at)` (dotted
+ * `freshness.window`, for resolveFreshness) and `tag(...at)` (a property
+ * collection — `queryMetadata { … }`, or the `materialization` envelope).
  */
-function fakeTag(
-   fields: Record<string, string> | undefined,
-   fs: FakeFreshnessSchedule | undefined,
-) {
+function fakeTagFromTree(tree: FakeTagTree) {
+   const walk = (at: string[]): string | FakeTagTree | undefined => {
+      let node: string | FakeTagTree | undefined = tree;
+      for (const segment of at) {
+         if (node === undefined || typeof node !== "object") return undefined;
+         node = node[segment];
+      }
+      return node;
+   };
    return {
-      *entries() {
-         for (const [key, value] of Object.entries(fields ?? {})) {
-            yield [key, { text: () => value }];
+      *entries(): Generator<[string, { text(): string | undefined }]> {
+         for (const [key, value] of Object.entries(tree)) {
+            yield [
+               key,
+               { text: () => (typeof value === "string" ? value : undefined) },
+            ];
          }
       },
       text(...at: string[]): string | undefined {
-         if (at.length === 1) {
-            return fields?.[at[0]];
-         }
-         if (at.length === 2 && at[0] === "freshness") {
-            return fs?.freshness?.[at[1] as "window" | "fallback"];
-         }
-         return undefined;
+         const node = walk(at);
+         return typeof node === "string" ? node : undefined;
+      },
+      tag(...at: string[]) {
+         const node = walk(at);
+         return node !== undefined && typeof node === "object"
+            ? fakeTagFromTree(node)
+            : undefined;
       },
    };
 }
 
+/** Assemble one tag layer's tree from the pieces a fake source declares. */
+function tagTree(pieces: {
+   fields?: Record<string, string>;
+   freshness?: FakeFreshnessSchedule;
+   queryMetadata?: Record<string, string>;
+   /** Contents of the `materialization` envelope (model-file layer). */
+   materialization?: {
+      freshness?: FakeFreshnessSchedule;
+      queryMetadata?: Record<string, string>;
+   };
+}): FakeTagTree {
+   const tree: FakeTagTree = { ...(pieces.fields ?? {}) };
+   if (pieces.freshness?.freshness) {
+      tree.freshness = { ...pieces.freshness.freshness } as FakeTagTree;
+   }
+   if (pieces.queryMetadata) tree.queryMetadata = { ...pieces.queryMetadata };
+   if (pieces.materialization) {
+      const envelope: FakeTagTree = {};
+      if (pieces.materialization.freshness?.freshness) {
+         envelope.freshness = {
+            ...pieces.materialization.freshness.freshness,
+         } as FakeTagTree;
+      }
+      if (pieces.materialization.queryMetadata) {
+         envelope.queryMetadata = { ...pieces.materialization.queryMetadata };
+      }
+      tree.materialization = envelope;
+   }
+   return tree;
+}
+
 export function fakeSource(opts: {
    name: string;
    sourceEntityId: string;
@@ -118,8 +162,17 @@ export function fakeSource(opts: {
    annotationFields?: Record<string, string>;
    /** Source-level (`#@`) freshness (dotted keys). */
    freshnessSchedule?: FakeFreshnessSchedule;
-   /** Model-file-level (`##`) freshness default. */
+   /** Model-file-level (`##`) freshness default, in the deprecated bare form. */
    modelFreshnessSchedule?: FakeFreshnessSchedule;
+   /** Model-file-level (`## materialization.*`) envelope declarations. */
+   modelMaterialization?: {
+      freshness?: FakeFreshnessSchedule;
+      queryMetadata?: Record<string, string>;
+   };
+   /** Source-level (`#@ persist queryMetadata.*`) per-query metadata. */
+   queryMetadata?: Record<string, string>;
+   /** Model-file-level bare (`## queryMetadata.*`) per-query metadata. */
+   modelQueryMetadata?: Record<string, string>;
    /**
     * Spy on the args Malloy's SQL generation is handed — chiefly the
     * `buildManifest` that resolves upstream persist references — so a test can
@@ -147,12 +200,24 @@ export function fakeSource(opts: {
       },
       annotations: {
          parseAsTag: () => ({
-            tag: fakeTag(fields, opts.freshnessSchedule),
+            tag: fakeTagFromTree(
+               tagTree({
+                  fields,
+                  freshness: opts.freshnessSchedule,
+                  queryMetadata: opts.queryMetadata,
+               }),
+            ),
          }),
       },
       modelAnnotations: {
          parseAsTag: () => ({
-            tag: fakeTag(undefined, opts.modelFreshnessSchedule),
+            tag: fakeTagFromTree(
+               tagTree({
+                  freshness: opts.modelFreshnessSchedule,
+                  queryMetadata: opts.modelQueryMetadata,
+                  materialization: opts.modelMaterialization,
+               }),
+            ),
          }),
       },
    } as unknown as PersistSource;