@milaboratories/pl-middle-layer 1.36.4 → 1.37.1

package/src/js_render/index.ts

@@ -2,6 +2,7 @@ import type { MiddleLayerEnvironment } from '../middle_layer/middle_layer';
 import type { Code, ConfigRenderLambda } from '@platforma-sdk/model';
 import type { ComputableRenderingOps } from '@milaboratories/computable';
 import { Computable } from '@milaboratories/computable';
+import type { QuickJSWASMModule } from 'quickjs-emscripten';
 import { Scope } from 'quickjs-emscripten';
 import type { DeadlineSettings } from './context';
 import { JsExecutionContext } from './context';
@@ -35,8 +36,9 @@ export function computableFromRF(
       return false;
     });
     const vm = scope.manage(runtime.newContext());
-    const rCtx = new JsExecutionContext(scope, vm, ctx, env,
-      (s) => { deadlineSettings = s; }, cCtx);
+    const rCtx = new JsExecutionContext(scope, vm,
+      (s) => { deadlineSettings = s; },
+      { computableCtx: cCtx, blockCtx: ctx, mlEnv: env });
 
     rCtx.evaluateBundle(code.content);
     const result = rCtx.runCallback(fh.handle);
@@ -49,7 +51,7 @@ export function computableFromRF(
       console.log(`Output ${fh.handle} scaffold calculated.`);
 
     return {
-      ir: rCtx.computablesToResolve,
+      ir: rCtx.computableHelper!.computablesToResolve,
       postprocessValue: (resolved: Record<string, unknown>, { unstableMarker, stable }) => {
         // resolving futures
         for (const [handle, value] of Object.entries(resolved)) rCtx.runCallback(handle, value);
@@ -75,3 +77,35 @@ export function computableFromRF(
     };
   }, ops);
 }
+
+export function executeSingleLambda(
+  quickJs: QuickJSWASMModule,
+  fh: ConfigRenderLambda,
+  code: Code,
+  ...args: unknown[]
+): unknown {
+  const scope = new Scope();
+  try {
+    const runtime = scope.manage(quickJs.newRuntime());
+    runtime.setMemoryLimit(1024 * 1024 * 8);
+    runtime.setMaxStackSize(1024 * 320);
+
+    let deadlineSettings: DeadlineSettings | undefined;
+    runtime.setInterruptHandler(() => {
+      if (deadlineSettings === undefined) return false;
+      if (Date.now() > deadlineSettings.deadline) return true;
+      return false;
+    });
+    const vm = scope.manage(runtime.newContext());
+    const rCtx = new JsExecutionContext(scope, vm,
+      (s) => { deadlineSettings = s; });
+
+    // Initializing the model
+    rCtx.evaluateBundle(code.content);
+
+    // Running the lambda
+    return rCtx.importObjectUniversal(rCtx.runCallback(fh.handle, ...args));
+  } finally {
+    scope.dispose();
+  }
+}

package/src/middle_layer/block.ts

@@ -43,7 +43,9 @@ export function blockArgsAndUiState(
   const uiState = ctx.uiState(cCtx);
   return {
     author: prj.getKeyValueAsJson<AuthorMarker>(blockArgsAuthorKey(blockId)),
+    // @TODO add deserialization caching
     args: deepFreeze(JSON.parse(ctx.args(cCtx))),
+    // @TODO add deserialization caching
     ui: uiState !== undefined ? deepFreeze(JSON.parse(uiState)) : undefined,
   };
 }

package/src/middle_layer/block_ctx.ts

@@ -13,6 +13,12 @@ import {
 import { allBlocks } from '../model/project_model_util';
 import { ResultPool } from '../pool/result_pool';
 
+export type BlockContextMaterialized = {
+  readonly blockId: string;
+  readonly args: string;
+  readonly uiState?: string;
+};
+
 export type BlockContextArgsOnly = {
   readonly blockId: string;
   readonly args: (cCtx: ComputableCtx) => string;

package/src/middle_layer/middle_layer.ts

@@ -37,6 +37,7 @@ import { V2RegistryProvider } from '../block_registry';
 import type { Dispatcher } from 'undici';
 import { RetryAgent } from 'undici';
 import { getDebugFlags } from '../debug';
+import { ProjectHelper } from '../model/project_helper';
 
 export interface MiddleLayerEnvironment {
   readonly pl: PlClient;
@@ -50,6 +51,7 @@ export interface MiddleLayerEnvironment {
   readonly blockUpdateWatcher: BlockUpdateWatcher;
   readonly quickJs: QuickJSWASMModule;
   readonly driverKit: MiddleLayerDriverKit;
+  readonly projectHelper: ProjectHelper;
 }
 
 /**
@@ -111,7 +113,7 @@ export class MiddleLayer {
     meta: ProjectMeta,
     author?: AuthorMarker,
   ): Promise<void> {
-    await withProjectAuthored(this.pl, rid, author, (prj) => {
+    await withProjectAuthored(this.env.projectHelper, this.pl, rid, author, (prj) => {
       prj.setMeta(meta);
     });
     await this.projectListTree.refreshState();
@@ -251,6 +253,8 @@ export class MiddleLayer {
       ops.frontendDownloadPath,
     );
 
+    const quickJs = await getQuickJS();
+
     const env: MiddleLayerEnvironment = {
       pl,
       signer: driverKit.signer,
@@ -266,7 +270,8 @@ export class MiddleLayer {
         http: retryHttpDispatcher,
         preferredUpdateChannel: ops.preferredUpdateChannel,
       }),
-      quickJs: await getQuickJS(),
+      quickJs,
+      projectHelper: new ProjectHelper(quickJs),
     };
 
     const openedProjects = new WatchableValue<ResourceId[]>([]);

package/src/middle_layer/project.ts

@@ -27,7 +27,7 @@ import { blockArgsAndUiState, blockOutputs } from './block';
 import type { FrontendData } from '../model/frontend';
 import type { ProjectStructure } from '../model/project_model';
 import { projectFieldName } from '../model/project_model';
-import { notEmpty } from '@milaboratories/ts-helpers';
+import { cachedDeserialize, notEmpty } from '@milaboratories/ts-helpers';
 import type { BlockPackInfo } from '../model/block_pack';
 import type {
   ProjectOverview,
@@ -107,7 +107,7 @@ export class Project {
   private async refreshLoop(): Promise<void> {
     while (!this.destroyed) {
       try {
-        await withProject(this.env.pl, this.rid, (prj) => {
+        await withProject(this.env.projectHelper, this.env.pl, this.rid, (prj) => {
           prj.doRefresh(this.env.ops.stagingRenderingRate);
         });
         await this.activeConfigs.getValue();
@@ -145,7 +145,7 @@ export class Project {
     const preparedBp = await this.env.bpPreparer.prepare(blockPackSpec);
     const blockCfgContainer = await this.env.bpPreparer.getBlockConfigContainer(blockPackSpec);
     const blockCfg = extractConfig(blockCfgContainer); // full content of this var should never be persisted
-    await withProjectAuthored(this.env.pl, this.rid, author, (mut) =>
+    await withProjectAuthored(this.env.projectHelper, this.env.pl, this.rid, author, (mut) =>
       mut.addBlock(
         {
           id: blockId,
@@ -177,7 +177,7 @@ export class Project {
   ): Promise<void> {
     const preparedBp = await this.env.bpPreparer.prepare(blockPackSpec);
     const blockCfg = extractConfig(await this.env.bpPreparer.getBlockConfigContainer(blockPackSpec));
-    await withProjectAuthored(this.env.pl, this.rid, author, (mut) =>
+    await withProjectAuthored(this.env.projectHelper, this.env.pl, this.rid, author, (mut) =>
       mut.migrateBlockPack(
         blockId,
         preparedBp,
@@ -189,7 +189,7 @@ export class Project {
 
   /** Deletes a block with all associated data. */
   public async deleteBlock(blockId: string, author?: AuthorMarker): Promise<void> {
-    await withProjectAuthored(this.env.pl, this.rid, author, (mut) => mut.deleteBlock(blockId));
+    await withProjectAuthored(this.env.projectHelper, this.env.pl, this.rid, author, (mut) => mut.deleteBlock(blockId));
     this.navigationStates.deleteBlock(blockId);
     await this.projectTree.refreshState();
   }
@@ -201,7 +201,7 @@ export class Project {
    * an error will be thrown instead.
    */
   public async reorderBlocks(blocks: string[], author?: AuthorMarker): Promise<void> {
-    await withProjectAuthored(this.env.pl, this.rid, author, (mut) => {
+    await withProjectAuthored(this.env.projectHelper, this.env.pl, this.rid, author, (mut) => {
       const currentStructure = mut.structure;
       if (currentStructure.groups.length !== 1)
         throw new Error('Unexpected project structure, non-sinular block group');
@@ -233,7 +233,7 @@ export class Project {
    * stale state.
    * */
   public async runBlock(blockId: string): Promise<void> {
-    await withProject(this.env.pl, this.rid, (mut) => mut.renderProduction([blockId], true));
+    await withProject(this.env.projectHelper, this.env.pl, this.rid, (mut) => mut.renderProduction([blockId], true));
     await this.projectTree.refreshState();
   }
 
@@ -243,7 +243,7 @@ export class Project {
    * calculated.
    * */
   public async stopBlock(blockId: string): Promise<void> {
-    await withProject(this.env.pl, this.rid, (mut) => mut.stopProduction(blockId));
+    await withProject(this.env.projectHelper, this.env.pl, this.rid, (mut) => mut.stopProduction(blockId));
     await this.projectTree.refreshState();
   }
 
@@ -262,7 +262,7 @@ export class Project {
    * in collaborative editing scenario.
    * */
   public async setBlockArgs(blockId: string, args: unknown, author?: AuthorMarker) {
-    await withProjectAuthored(this.env.pl, this.rid, author, (mut) =>
+    await withProjectAuthored(this.env.projectHelper, this.env.pl, this.rid, author, (mut) =>
       mut.setArgs([{ blockId, args: canonicalize(args)! }]),
     );
     await this.projectTree.refreshState();
@@ -275,7 +275,7 @@ export class Project {
    * in collaborative editing scenario.
    * */
   public async setUiState(blockId: string, uiState: unknown, author?: AuthorMarker) {
-    await withProjectAuthored(this.env.pl, this.rid, author, (mut) =>
+    await withProjectAuthored(this.env.projectHelper, this.env.pl, this.rid, author, (mut) =>
       mut.setUiState(blockId, uiState === undefined ? undefined : canonicalize(uiState)!),
     );
     await this.projectTree.refreshState();
@@ -301,7 +301,7 @@ export class Project {
     uiState: unknown,
     author?: AuthorMarker,
   ) {
-    await withProjectAuthored(this.env.pl, this.rid, author, (mut) => {
+    await withProjectAuthored(this.env.projectHelper, this.env.pl, this.rid, author, (mut) => {
       mut.setArgs([{ blockId, args: canonicalize(args)! }]);
       mut.setUiState(blockId, canonicalize(uiState));
     });
@@ -310,7 +310,7 @@ export class Project {
 
   /** Update block settings */
   public async setBlockSettings(blockId: string, newValue: BlockSettings) {
-    await withProjectAuthored(this.env.pl, this.rid, undefined, (mut) => {
+    await withProjectAuthored(this.env.projectHelper, this.env.pl, this.rid, undefined, (mut) => {
       mut.setBlockSettings(blockId, newValue);
     });
     await this.projectTree.refreshState();
@@ -327,10 +327,8 @@ export class Project {
         (await tx.getField(field(bpHolderRid, Pl.HolderRefField))).value,
       );
       const bpData = await tx.getResourceData(bpRid, false);
-      const config = extractConfig((JSON.parse(
-        Buffer.from(notEmpty(bpData.data)).toString('utf-8'),
-      ) as BlockPackInfo).config);
-      await withProjectAuthored(tx, this.rid, author, (prj) => {
+      const config = extractConfig((cachedDeserialize(notEmpty(bpData.data)) as BlockPackInfo).config);
+      await withProjectAuthored(this.env.projectHelper, tx, this.rid, author, (prj) => {
         prj.setArgs([{ blockId, args: canonicalize(config.initialArgs)! }]);
         prj.setUiState(blockId, canonicalize(config.initialUiState));
       });
@@ -434,7 +432,7 @@ export class Project {
 
   public static async init(env: MiddleLayerEnvironment, rid: ResourceId): Promise<Project> {
     // Doing a no-op mutation to apply all migration and schema fixes
-    await withProject(env.pl, rid, (_) => {});
+    await withProject(env.projectHelper, env.pl, rid, (_) => {});
 
     // Loading project tree
     const projectTree = await SynchronizedTreeState.init(

package/src/middle_layer/project_overview.ts

@@ -29,11 +29,12 @@ import type { BlockSection } from '@platforma-sdk/model';
 import { computableFromCfgOrRF } from './render';
 import type { NavigationStates } from './navigation_states';
 import { getBlockPackInfo } from './util';
-import { resourceIdToString } from '@milaboratories/pl-client';
+import { resourceIdToString, type ResourceId } from '@milaboratories/pl-client';
 import * as R from 'remeda';
 
 type BlockInfo = {
-  currentArguments: Record<string, unknown>;
+  argsRid: ResourceId;
+  currentArguments: unknown;
   prod?: ProdState;
 };
 
@@ -127,10 +128,18 @@ export function projectOverview(
           };
         }
 
-        infos.set(id, { currentArguments, prod });
+        infos.set(id, { currentArguments, prod, argsRid: cInputs.resourceInfo.id });
       }
 
-      const currentGraph = productionGraph(structure, (id) => infos.get(id)!.currentArguments);
+      const currentGraph = productionGraph(structure, (id) => {
+        const bpInfo = getBlockPackInfo(prj, id)!;
+        const bInfo = infos.get(id)!;
+        const args = bInfo.currentArguments;
+        return {
+          args,
+          enrichmentTargets: env.projectHelper.getEnrichmentTargets(() => bpInfo.cfg, () => args, { argsRid: bInfo.argsRid, blockPackRid: bpInfo.bpResourceId }),
+        };
+      });
 
       const limbo = new Set(renderingState.blocksInLimbo);
 

package/src/middle_layer/util.ts

@@ -1,5 +1,6 @@
 import type { PlTreeNodeAccessor } from '@milaboratories/pl-tree';
 import { projectFieldName } from '../model/project_model';
+import type { ResourceId } from '@milaboratories/pl-client';
 import { Pl } from '@milaboratories/pl-client';
 import { ifNotUndef } from '../cfg_render/util';
 import type { BlockPackInfo } from '../model/block_pack';
@@ -7,6 +8,7 @@ import type { BlockConfig } from '@platforma-sdk/model';
 import { extractConfig } from '@platforma-sdk/model';
 
 export type BlockPackInfoAndId = {
+  readonly bpResourceId: ResourceId;
   /** To be added to computable keys, to force reload on config change */
   readonly bpId: string;
   /** Full block-pack info */
@@ -32,7 +34,7 @@ export function getBlockPackInfo(
     (bpAcc) => {
       const info = bpAcc.getDataAsJson<BlockPackInfo>()!;
       const cfg = extractConfig(info.config);
-      return { bpId: bpAcc.resourceInfo.id.toString(), info, cfg };
+      return { bpResourceId: bpAcc.resourceInfo.id, bpId: bpAcc.resourceInfo.id.toString(), info, cfg };
     },
   );
 }

package/src/model/args.ts

@@ -1,8 +1,9 @@
 import { assertNever } from '@milaboratories/ts-helpers';
 import type { PlRef } from '@platforma-sdk/model';
 
-export function outputRef(blockId: string, name: string): PlRef {
-  return { __isRef: true, blockId, name };
+export function outputRef(blockId: string, name: string, requireEnrichments?: boolean): PlRef {
+  if (requireEnrichments) return { __isRef: true, blockId, name, requireEnrichments };
+  else return { __isRef: true, blockId, name };
 }
 
 export function isBlockOutputReference(obj: unknown): obj is PlRef {
@@ -32,8 +33,11 @@ function addAllReferencedBlocks(result: BlockUpstreams, node: unknown, allowed?:
       if (node === null) return;
 
       if (isBlockOutputReference(node)) {
-        if (allowed === undefined || allowed.has(node.blockId)) result.upstreams.add(node.blockId);
-        else result.missingReferences = true;
+        if (allowed === undefined || allowed.has(node.blockId)) {
+          result.upstreams.add(node.blockId);
+          if (node.requireEnrichments)
+            result.upstreamsRequiringEnrichments.add(node.blockId);
+        } else result.missingReferences = true;
       } else if (Array.isArray(node)) {
         for (const child of node) addAllReferencedBlocks(result, child, allowed);
       } else {
@@ -49,15 +53,17 @@ function addAllReferencedBlocks(result: BlockUpstreams, node: unknown, allowed?:
 }
 
 export interface BlockUpstreams {
-  /** Direct block dependencies */
+  /** All direct block dependencies */
   upstreams: Set<string>;
+  /** Direct block dependencies which enrichments are also required by current block */
+  upstreamsRequiringEnrichments: Set<string>;
   /** True if not-allowed references was encountered */
   missingReferences: boolean;
 }
 
 /** Extracts all resource ids referenced by args object. */
 export function inferAllReferencedBlocks(args: unknown, allowed?: Set<string>): BlockUpstreams {
-  const result = { upstreams: new Set<string>(), missingReferences: false };
+  const result = { upstreams: new Set<string>(), upstreamsRequiringEnrichments: new Set<string>(), missingReferences: false };
   addAllReferencedBlocks(result, args, allowed);
   return result;
 }

package/src/model/project_helper.ts

@@ -0,0 +1,41 @@
+import type { BlockConfig, PlRef } from '@platforma-sdk/model';
+import { LRUCache } from 'lru-cache';
+import type { QuickJSWASMModule } from 'quickjs-emscripten';
+import { executeSingleLambda } from '../js_render';
+import type { ResourceId } from '@milaboratories/pl-client';
+
+type EnrichmentTargetsRequest = {
+  blockConfig: () => BlockConfig;
+  args: () => unknown;
+};
+
+type EnrichmentTargetsValue = {
+  value: PlRef[] | undefined;
+};
+
+export class ProjectHelper {
+  private readonly enrichmentTargetsCache = new LRUCache<string, EnrichmentTargetsValue, EnrichmentTargetsRequest>({
+    max: 256,
+    memoMethod: (_key, _value, { context }) => {
+      return { value: this.calculateEnrichmentTargets(context) };
+    },
+  });
+
+  constructor(private readonly quickJs: QuickJSWASMModule) {}
+
+  private calculateEnrichmentTargets(req: EnrichmentTargetsRequest): PlRef[] | undefined {
+    const blockConfig = req.blockConfig();
+    if (blockConfig.enrichmentTargets === undefined) return undefined;
+    const args = req.args();
+    const result = executeSingleLambda(this.quickJs, blockConfig.enrichmentTargets, blockConfig.code!, args) as PlRef[];
+    return result;
+  }
+
+  public getEnrichmentTargets(blockConfig: () => BlockConfig, args: () => unknown, key?: { argsRid: ResourceId; blockPackRid: ResourceId }): PlRef[] | undefined {
+    const req = { blockConfig, args };
+    if (key === undefined)
+      return this.calculateEnrichmentTargets(req);
+    const cacheKey = `${key.argsRid}:${key.blockPackRid}`;
+    return this.enrichmentTargetsCache.memo(cacheKey, { context: req }).value;
+  }
+}

package/src/model/project_model_util.test.ts

@@ -4,7 +4,7 @@ import { outputRef } from './args';
 import { PlRef } from '@platforma-sdk/model';
 
 function toRefs(...ids: string[]): PlRef[] {
-  return ids.map((id) => outputRef(id, ''));
+  return ids.map((id) => outputRef(id, '', true));
 }
 
 function simpleStructure(...ids: string[]): ProjectStructure {
@@ -26,7 +26,10 @@ describe('simple traverse', () => {
   inputs.set('b4', toRefs('b3'));
   inputs.set('b5', toRefs('b4'));
   inputs.set('b6', toRefs('b2', 'b4'));
-  const pGraph1 = productionGraph(struct1, (id) => inputs.get(id) ?? null);
+  const pGraph1 = productionGraph(struct1, (id) => ({
+    args: inputs.get(id) ?? null,
+    enrichmentTargets: undefined,
+  }));
 
   test.each([
     { roots: ['b2'], expectedDirectUpstreams: ['b1'], expectedDirectDownstreams: ['b6'], expectedUpstreams: ['b1'], expectedDownstreams: ['b6'] },
@@ -51,7 +54,10 @@ describe('simple diff', () => {
   inputs.set('b2', toRefs('b1'));
   inputs.set('b4', toRefs('b3'));
   const sGraph1 = stagingGraph(struct1);
-  const pGraph1 = productionGraph(struct1, (id) => inputs.get(id) ?? null);
+  const pGraph1 = productionGraph(struct1, (id) => ({
+    args: inputs.get(id) ?? null,
+    enrichmentTargets: undefined,
+  }));
 
   test.each([
     { struct2a: ['b1', 'b2', 'b3', 'b4'], expectedS: [], expectedP: [] },
@@ -79,7 +85,10 @@ describe('simple diff', () => {
   ])('$struct2a', ({ struct2a, expectedS, expectedP, onlyA, onlyB }) => {
     const struct2: ProjectStructure = simpleStructure(...struct2a);
     const sGraph2 = stagingGraph(struct2);
-    const pGraph2 = productionGraph(struct2, (id) => inputs.get(id) ?? null);
+    const pGraph2 = productionGraph(struct2, (id) => ({
+      args: inputs.get(id) ?? null,
+      enrichmentTargets: undefined,
+    }));
     const sDiff = graphDiff(sGraph1, sGraph2);
     const pDiff = graphDiff(pGraph1, pGraph2);
     expect(sDiff.onlyInA).toEqual(new Set(onlyA ?? []));

package/src/model/project_model_util.ts

@@ -1,6 +1,7 @@
 import type { Block, ProjectStructure } from './project_model';
 import type { Optional, Writable } from 'utility-types';
 import { inferAllReferencedBlocks } from './args';
+import type { PlRef } from '@milaboratories/pl-model-common';
 
 export function allBlocks(structure: ProjectStructure): Iterable<Block> {
   return {
@@ -19,10 +20,16 @@ export interface BlockGraphNode {
   readonly directUpstream: Set<string>;
   /** Downstreams, calculated accounting for potential indirect block dependencies  */
   readonly downstream: Set<string>;
+  /** Downstream blocks enriching current block's outputs */
+  readonly enrichments: Set<string>;
   /** Direct downstreams listing our block in it's args */
   readonly directDownstream: Set<string>;
+  /** Upstream blocks that current block may enrich with its exports */
+  readonly enrichmentTargets: Set<string>;
 }
 
+export type BlockGraphDirection = 'upstream' | 'downstream' | 'directUpstream' | 'directDownstream' | 'enrichments' | 'enrichmentTargets';
+
 export class BlockGraph {
   /** Nodes are stored in the map in topological order */
   public readonly nodes: Map<string, BlockGraphNode>;
@@ -32,7 +39,7 @@ export class BlockGraph {
   }
 
   public traverseIds(
-    direction: 'upstream' | 'downstream' | 'directUpstream' | 'directDownstream',
+    direction: BlockGraphDirection,
     ...rootBlockIds: string[]
   ): Set<string> {
     const all = new Set<string>();
@@ -41,7 +48,7 @@ export class BlockGraph {
   }
 
   public traverseIdsExcludingRoots(
-    direction: 'upstream' | 'downstream' | 'directUpstream' | 'directDownstream',
+    direction: BlockGraphDirection,
     ...rootBlockIds: string[]
   ): Set<string> {
     const result = this.traverseIds(direction, ...rootBlockIds);
@@ -50,7 +57,7 @@ export class BlockGraph {
   }
 
   public traverse(
-    direction: 'upstream' | 'downstream' | 'directUpstream' | 'directDownstream',
+    direction: BlockGraphDirection,
     rootBlockIds: string[],
     cb: (node: BlockGraphNode) => void,
   ): void {
@@ -75,7 +82,7 @@ export class BlockGraph {
 }
 
 export function stagingGraph(structure: ProjectStructure) {
-  type WNode = Optional<Writable<BlockGraphNode>, 'upstream' | 'downstream' | 'directUpstream' | 'directDownstream'>;
+  type WNode = Optional<Writable<BlockGraphNode>, BlockGraphDirection>;
   const result = new Map<string, WNode>();
 
   // Simple dependency graph from previous to next
@@ -92,9 +99,11 @@ export function stagingGraph(structure: ProjectStructure) {
     result.set(id, current);
     if (previous === undefined) {
       current.directUpstream = current.upstream = new Set<string>();
+      current.enrichments = current.enrichmentTargets = new Set<string>();
     } else {
       current.directUpstream = current.upstream = new Set<string>([previous.id]);
       previous.directDownstream = previous.downstream = new Set<string>([current.id]);
+      previous.enrichments = previous.enrichmentTargets = new Set<string>();
     }
 
     previous = current;
@@ -104,9 +113,14 @@ export function stagingGraph(structure: ProjectStructure) {
   return new BlockGraph(result as Map<string, BlockGraphNode>);
 }
 
+export type ProductionGraphBlockInfo = {
+  args: unknown;
+  enrichmentTargets: PlRef[] | undefined;
+};
+
 export function productionGraph(
   structure: ProjectStructure,
-  argsProvider: (blockId: string) => unknown,
+  infoProvider: (blockId: string) => ProductionGraphBlockInfo | undefined,
 ): BlockGraph {
   const resultMap = new Map<string, BlockGraphNode>();
   // result graph is constructed to be able to perform traversal on incomplete graph
@@ -118,16 +132,19 @@ export function productionGraph(
   // those dependencies that are possible under current topology
   const allAbove = new Set<string>();
   for (const { id } of allBlocks(structure)) {
-    const args = argsProvider(id);
+    const info = infoProvider(id);
 
     // skipping those blocks for which we don't have args
-    if (args === undefined) continue;
+    if (info === undefined) continue;
 
-    const directUpstreams = inferAllReferencedBlocks(args, allAbove);
+    const references = inferAllReferencedBlocks(info.args, allAbove);
 
     // The algorithm here adds all downstream blocks of direct upstreams as potential upstreams.
     // They may produce additional columns, anchored in our direct upstream, those columns might be needed by the workflow.
-    const potentialUpstreams = resultGraph.traverseIds('directDownstream', ...directUpstreams.upstreams);
+    const potentialUpstreams = new Set([
+      ...references.upstreams,
+      ...resultGraph.traverseIds('enrichments', ...references.upstreamsRequiringEnrichments),
+    ]);
 
     // To minimize complexity of the graph, we leave only the closest upstreams, removing all their transitive dependencies,
     // relying on the traversal mechanisms in BContexts and on UI level to connect our block with all upstreams
@@ -141,17 +158,25 @@ export function productionGraph(
       }
     }
 
+    // default assumption is that all direct upstreams are enrichment targets
+    const enrichmentTargets = info.enrichmentTargets === undefined
+      ? new Set(references.upstreams)
+      : new Set(info.enrichmentTargets.map((t) => t.blockId));
+
     const node: BlockGraphNode = {
       id,
-      missingReferences: directUpstreams.missingReferences,
+      missingReferences: references.missingReferences,
       upstream: upstreams,
-      directUpstream: directUpstreams.upstreams,
+      directUpstream: references.upstreams,
+      enrichmentTargets,
       downstream: new Set<string>(), // will be populated from downstream blocks
       directDownstream: new Set<string>(), // will be populated from downstream blocks
+      enrichments: new Set<string>(), // will be populated from downstream blocks
     };
     resultMap.set(id, node);
-    directUpstreams.upstreams.forEach((dep) => resultMap.get(dep)!.directDownstream.add(id));
+    references.upstreams.forEach((dep) => resultMap.get(dep)!.directDownstream.add(id));
     upstreams.forEach((dep) => resultMap.get(dep)!.downstream.add(id));
+    enrichmentTargets.forEach((dep) => resultMap.get(dep)?.enrichments.add(id));
     allAbove.add(id);
   }