Support ref() and DAG dependencies in PropertyGraph

core/actions/property_graph.ts

@@ -1,15 +1,32 @@
 import { verifyObjectMatchesProto, VerifyProtoErrorBehaviour } from "df/common/protos";
 import { ActionBuilder } from "df/core/actions";
+import { Declaration } from "df/core/actions/declaration";
 import { Session } from "df/core/session";
+import { checkAssertionsForDependency } from "df/core/utils";
 import { dataform } from "df/protos/ts";
 
 const CATALOG_NOT_SUPPORTED_MESSAGE =
   "Catalog-based data sources (BigLake/Iceberg/external catalogs) are not yet supported.";
 
+function entityRefKey(name: string): string {
+  return `entity:${name}`;
+}
+
+function relationshipRefKey(name: string): string {
+  return `relationship:${name}`;
+}
+
 export class PropertyGraph extends ActionBuilder<dataform.PropertyGraph> {
   public session: Session;
+  public dependOnDependencyAssertions: boolean = false;
 
   private proto = dataform.PropertyGraph.create({ description: "", disabled: false });
+  // Populated during construction for entity/relationship configs that use a ref-style
+  // dataSource; drained in finalize() once all actions are indexed, since the target
+  // schema/database of a ref cannot be resolved until every action is known.
+  // Keys are produced by entityRefKey / relationshipRefKey.
+  private pendingRefs = new Map<string, dataform.ITarget>();
+  private dependencyKeys = new Set<string>();
 
   constructor(session?: Session, unverifiedConfig?: any, filename?: string) {
     super(session);
@@ -61,8 +78,6 @@ export class PropertyGraph extends ActionBuilder<dataform.PropertyGraph> {
     this.validateUniqueNames(config.name);
     this.resolveEndpointDefaults();
     this.validateEndpointReferences(config.name);
-
-    this.proto.graphBody = this.emitGraphBody();
   }
 
   public getFileName() {
@@ -74,20 +89,66 @@ export class PropertyGraph extends ActionBuilder<dataform.PropertyGraph> {
   }
 
   public compile() {
-    return verifyObjectMatchesProto(
+    this.proto = verifyObjectMatchesProto(
       dataform.PropertyGraph,
       this.proto,
       VerifyProtoErrorBehaviour.SUGGEST_REPORTING_TO_DATAFORM_TEAM
     );
+    return this.proto;
+  }
+
+  public finalize() {
+    let allResolved = true;
+    for (const entity of this.proto.entities) {
+      if (!this.resolveRefIntoDataSource(entity, entityRefKey(entity.name))) {
+        allResolved = false;
+      }
+    }
+    for (const rel of this.proto.relationships) {
+      if (!this.resolveRefIntoDataSource(rel, relationshipRefKey(rel.name))) {
+        allResolved = false;
+      }
+    }
+    if (allResolved) {
+      this.proto.graphBody = this.emitGraphBody();
+    }
+  }
+
+  private resolveRefIntoDataSource(
+    entityOrRel: dataform.IGraphEntity | dataform.IGraphRelationship,
+    key: string
+  ): boolean {
+    const rawRef = this.pendingRefs.get(key);
+    if (!rawRef) {
+      return true;
+    }
+    const resolved = this.session.indexedActions.find(rawRef);
+    if (resolved.length !== 1) {
+      return false;
+    }
+    const resolvedAction = resolved[0];
+    const target = resolvedAction.getTarget();
+    if (resolvedAction instanceof Declaration) {
+      entityOrRel.dataSource = target;
+    } else {
+      entityOrRel.dataSource = dataform.Target.create({
+        database: target.database && this.session.finalizeDatabase(target.database),
+        schema: this.session.finalizeSchema(target.schema),
+        name: this.session.finalizeName(target.name)
+      });
+    }
+    return true;
   }
 
   private normalizeEntitiesAndRelationships(unverifiedConfig: any) {
     for (const entity of unverifiedConfig.entities || []) {
+      normalizeRef(entity);
       normalizeKeys(entity);
       normalizeFields(entity);
       normalizeFieldsOnLabels(entity);
     }
     for (const relationship of unverifiedConfig.relationships || []) {
+      normalizeRef(relationship);
       normalizeKeys(relationship);
       normalizeFields(relationship);
       normalizeFieldsOnLabels(relationship);
@@ -113,7 +174,7 @@ export class PropertyGraph extends ActionBuilder<dataform.PropertyGraph> {
     }
     const entityName = entityConfig.name;
     const where = `Property graph '${graphName}': entity '${entityName}'`;
-    const dataSource = this.resolveDataSource(entityConfig, where);
+    const dataSource = this.resolveDataSource(entityConfig, entityRefKey(entityName), where);
     if (!entityConfig.keys || entityConfig.keys.length === 0) {
       throw new Error(`${where} must declare 'keys'.`);
     }
@@ -137,7 +198,7 @@ export class PropertyGraph extends ActionBuilder<dataform.PropertyGraph> {
     }
     const relName = relConfig.name;
     const where = `Property graph '${graphName}': relationship '${relName}'`;
-    const dataSource = this.resolveDataSource(relConfig, where);
+    const dataSource = this.resolveDataSource(relConfig, relationshipRefKey(relName), where);
 
     if (!relConfig.source) {
       throw new Error(`${where} must declare 'source'.`);
@@ -163,6 +224,7 @@ export class PropertyGraph extends ActionBuilder<dataform.PropertyGraph> {
 
   private resolveDataSource(
     entityOrRel: dataform.IGraphEntityConfig | dataform.IGraphRelationshipConfig,
+    pendingKey: string,
     where: string
   ): dataform.Target {
     if (entityOrRel.dataSourceCatalog) {
@@ -187,6 +249,29 @@ export class PropertyGraph extends ActionBuilder<dataform.PropertyGraph> {
       }
       return dataform.Target.create({ name: ds.table, schema: dataset, database: project });
     }
+    if (entityOrRel.dataSourceRef) {
+      const ref = entityOrRel.dataSourceRef;
+      if (!ref.name) {
+        throw new Error(`${where}: 'ref' must include a 'name'.`);
+      }
+      const rawRef: dataform.ITarget = { name: ref.name };
+      if (ref.schema) {
+        rawRef.schema = ref.schema;
+      }
+      if (ref.database) {
+        rawRef.database = ref.database;
+      }
+      this.pendingRefs.set(pendingKey, rawRef);
+      const depKey = `${rawRef.database || ""}.${rawRef.schema || ""}.${rawRef.name}`;
+      if (!this.dependencyKeys.has(depKey)) {
+        this.dependencyKeys.add(depKey);
+        const depTarget = checkAssertionsForDependency(this, rawRef);
+        if (depTarget) {
+          this.proto.dependencyTargets.push(depTarget);
+        }
+      }
+      return undefined;
+    }
     throw new Error(`${where}: must declare a data source.`);
   }
 
@@ -268,6 +353,21 @@ export class PropertyGraph extends ActionBuilder<dataform.PropertyGraph> {
   }
 }
 
+function normalizeRef(entityOrRel: any) {
+  if (entityOrRel.ref === undefined || entityOrRel.ref === null) {
+    return;
+  }
+  if (typeof entityOrRel.ref === "string") {
+    entityOrRel.dataSourceRef = { name: entityOrRel.ref };
+    delete entityOrRel.ref;
+    return;
+  }
+  if (typeof entityOrRel.ref === "object") {
+    entityOrRel.dataSourceRef = entityOrRel.ref;
+    delete entityOrRel.ref;
+  }
+}
+
 function normalizeKeys(entityOrRel: any) {
   if (typeof entityOrRel.keys === "string") {
     entityOrRel.keys = [entityOrRel.keys];

core/actions/property_graph_test.ts

@@ -16,7 +16,14 @@ function makeSession(): Session {
 }
 
 function compile(config: any, filename = "definitions/graph.yaml"): dataform.PropertyGraph {
-  return new PropertyGraph(makeSession(), config, filename).compile();
+  const action = new PropertyGraph(makeSession(), config, filename);
+  const compiled = action.compile();
+  action.finalize();
+  return compiled;
+}
+
+function build(config: any, filename = "definitions/graph.yaml"): PropertyGraph {
+  return new PropertyGraph(makeSession(), config, filename);
 }
 
 const graphTarget = (name: string) => ({
@@ -1352,4 +1359,138 @@ suite("property_graph", () => {
       })
     );
   });
+
+  test("scalar ref normalizes to dataSourceRef and populates dependencyTargets", () => {
+    const compiled = build({
+      name: "G",
+      entities: [
+        {
+          name: "A",
+          ref: "books",
+          keys: ["id"]
+        }
+      ]
+    }).compile();
+
+    expect(asPlainObject(compiled)).deep.equals(
+      asPlainObject({
+        target: graphTarget("G"),
+        canonicalTarget: graphTarget("G"),
+        dependencyTargets: [{ name: "books", includeDependentAssertions: false }],
+        fileName: "definitions/graph.yaml",
+        description: "",
+        disabled: false,
+        entities: [{ name: "A", keys: ["id"] }]
+      })
+    );
+  });
+
+  test("object ref preserves schema and database in dependencyTargets", () => {
+    const compiled = build({
+      name: "G",
+      entities: [
+        {
+          name: "A",
+          ref: { name: "books", schema: "analytics", database: "proj" },
+          keys: ["id"]
+        }
+      ]
+    }).compile();
+
+    expect(asPlainObject(compiled)).deep.equals(
+      asPlainObject({
+        target: graphTarget("G"),
+        canonicalTarget: graphTarget("G"),
+        dependencyTargets: [
+          { database: "proj", schema: "analytics", name: "books", includeDependentAssertions: false }
+        ],
+        fileName: "definitions/graph.yaml",
+        description: "",
+        disabled: false,
+        entities: [{ name: "A", keys: ["id"] }]
+      })
+    );
+  });
+
+  test("duplicate refs across entities dedupe in dependencyTargets", () => {
+    const compiled = build({
+      name: "G",
+      entities: [
+        { name: "A", ref: "books", keys: ["id"] },
+        { name: "B", ref: "books", keys: ["id"] }
+      ]
+    }).compile();
+
+    expect(asPlainObject(compiled)).deep.equals(
+      asPlainObject({
+        target: graphTarget("G"),
+        canonicalTarget: graphTarget("G"),
+        dependencyTargets: [{ name: "books", includeDependentAssertions: false }],
+        fileName: "definitions/graph.yaml",
+        description: "",
+        disabled: false,
+        entities: [
+          { name: "A", keys: ["id"] },
+          { name: "B", keys: ["id"] }
+        ]
+      })
+    );
+  });
+
+  test("ref on relationship also normalizes and populates dependencyTargets", () => {
+    const compiled = build({
+      name: "G",
+      entities: [
+        { name: "A", dataSourceString: "p.d.A", keys: ["id"] },
+        { name: "B", dataSourceString: "p.d.B", keys: ["id"] }
+      ],
+      relationships: [
+        {
+          name: "R",
+          ref: "wrote",
+          keys: ["a_id", "b_id"],
+          source: { entity: "A", joinKeys: ["a_id"] },
+          destination: { entity: "B", joinKeys: ["b_id"] }
+        }
+      ]
+    }).compile();
+
+    expect(asPlainObject(compiled)).deep.equals(
+      asPlainObject({
+        target: graphTarget("G"),
+        canonicalTarget: graphTarget("G"),
+        dependencyTargets: [{ name: "wrote", includeDependentAssertions: false }],
+        fileName: "definitions/graph.yaml",
+        description: "",
+        disabled: false,
+        entities: [
+          { name: "A", dataSource: { database: "p", schema: "d", name: "A" }, keys: ["id"] },
+          { name: "B", dataSource: { database: "p", schema: "d", name: "B" }, keys: ["id"] }
+        ],
+        relationships: [
+          {
+            name: "R",
+            keys: ["a_id", "b_id"],
+            source: { entity: "A", relationshipColumns: ["a_id"], entityColumns: ["id"] },
+            destination: { entity: "B", relationshipColumns: ["b_id"], entityColumns: ["id"] }
+          }
+        ]
+      })
+    );
+  });
+
+  test("errors when ref has no name", () => {
+    expect(() =>
+      compile({
+        name: "G",
+        entities: [
+          {
+            name: "A",
+            ref: { schema: "s" },
+            keys: ["id"]
+          }
+        ]
+      })
+    ).to.throw("'ref' must include a 'name'");
+  });
 });