@malloy-publisher/server 0.0.233 → 0.0.235

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;

package/src/service/model.spec.ts

@@ -895,12 +895,14 @@ describe("service/model", () => {
    describe("runtime storage failure → freshnessFallback=live", () => {
       const originalMode = process.env.PERSIST_STORAGE_MODE;
       const originalDefaultRows = process.env.PUBLISHER_DEFAULT_QUERY_ROW_LIMIT;
+      const originalMetadata = process.env.PUBLISHER_QUERY_METADATA;
 
       afterEach(() => {
          sinon.restore();
          for (const [name, original] of [
             ["PERSIST_STORAGE_MODE", originalMode],
             ["PUBLISHER_DEFAULT_QUERY_ROW_LIMIT", originalDefaultRows],
+            ["PUBLISHER_QUERY_METADATA", originalMetadata],
          ] as const) {
             if (original === undefined) delete process.env[name];
             else process.env[name] = original;
@@ -936,7 +938,10 @@ describe("service/model", () => {
             getPreparedResult:
                opts.storageFailsAt === "prepare"
                   ? sinon.stub().rejects(storageErr)
-                  : sinon.stub().resolves({ resultExplore: { limit: 0 } }),
+                  : sinon.stub().resolves({
+                       resultExplore: { limit: 0 },
+                       connectionName: "lake",
+                    }),
             run:
                opts.storageFailsAt === "run"
                   ? sinon.stub().rejects(storageErr)
@@ -954,6 +959,7 @@ describe("service/model", () => {
          const liveRunnable = {
             getPreparedResult: sinon.stub().resolves({
                resultExplore: { limit: opts.livePreparedLimit ?? 0 },
+               connectionName: "live_pg",
             }),
             run: liveRun,
          };
@@ -1020,6 +1026,44 @@ describe("service/model", () => {
          expect(liveRun.firstCall.args[0].rowLimit).toBe(250);
       });
 
+      it("tags the live retry, with the LIVE connection's layers and the same id", async () => {
+         // Two failures in one: the statement that actually answered the query
+         // carried no attribution at all, so the full-cost fallback — the one an
+         // operator goes looking for on a bill — landed in the untagged bucket;
+         // and the response still returned the id from the bag resolved before
+         // the storage attempt, pointing a caller at the statement that FAILED.
+         // The layers must come from the live connection: on this tier the store
+         // is routinely a different connection, with different enforced
+         // properties.
+         process.env.PUBLISHER_QUERY_METADATA = "on";
+         const { model, liveRun } = routedModel({
+            shapeBindings: [binding("daily", "live")],
+            storageFailsAt: "run",
+         });
+
+         const result = await model.getQueryResults(
+            undefined,
+            undefined,
+            "run: daily -> x",
+            undefined,
+            undefined,
+            undefined,
+            undefined,
+            {
+               correlationId: "corr-1",
+               connectionMetadata: (connectionName: string) =>
+                  connectionName === "live_pg"
+                     ? { default: null, enforced: { tenant: "acme" } }
+                     : { default: null, enforced: { tenant: "wrong" } },
+            },
+         );
+
+         const attached = liveRun.firstCall.args[0].queryMetadata;
+         expect(attached.tenant).toBe("acme");
+         expect(attached.query_id).toBe("corr-1");
+         expect(result.queryCorrelationId).toBe("corr-1");
+      });
+
       it("does not record a live-served answer as a storage hit", async () => {
          // The hit rate is the tier's headline KPI; counting a fallback as a hit
          // makes it RISE while the tier is broken.