@nicia-ai/typegraph 0.9.1 → 0.10.0

package/dist/index.cjs

@@ -2522,14 +2522,8 @@ function compileAggregateExprFromSource(expr, dialect) {
   const { field: field2 } = expr;
   const fn = expr.function;
   switch (fn) {
-    case "count": {
-      const cteAlias = `cte_${field2.alias}`;
-      return drizzleOrm.sql`COUNT(${drizzleOrm.sql.raw(cteAlias)}.${drizzleOrm.sql.raw(field2.alias)}_id)`;
-    }
-    case "countDistinct": {
-      const cteAlias = `cte_${field2.alias}`;
-      return drizzleOrm.sql`COUNT(DISTINCT ${drizzleOrm.sql.raw(cteAlias)}.${drizzleOrm.sql.raw(field2.alias)}_id)`;
-    }
+    case "count":
+    case "countDistinct":
     case "sum":
     case "avg":
     case "min":
@@ -2541,6 +2535,9 @@ function compileAggregateExprFromSource(expr, dialect) {
         field2.valueType,
         cteAlias
       );
+      if (fn === "countDistinct") {
+        return drizzleOrm.sql`COUNT(DISTINCT ${column})`;
+      }
       return drizzleOrm.sql`${drizzleOrm.sql.raw(fn.toUpperCase())}(${column})`;
     }
     default: {
@@ -3036,7 +3033,7 @@ function compileRecursiveCte(ast, traversal, graphId, ctx, requiredColumnsByAlia
     recursiveBaseWhereClauses.push(cycleCheck);
   }
   recursiveBaseWhereClauses.push(...edgePredicates, ...targetNodePredicates);
-  function compileRecursiveBranch(branch) {
+  function compileRecursiveBranch2(branch) {
     const recursiveFilterClauses = [
       ...recursiveBaseWhereClauses,
       compileKindFilter2(branch.edgeKinds, "e.kind"),
@@ -3091,7 +3088,7 @@ function compileRecursiveCte(ast, traversal, graphId, ctx, requiredColumnsByAlia
   const directTargetField = direction === "out" ? "to_id" : "from_id";
   const directJoinKindField = direction === "out" ? "from_kind" : "to_kind";
   const directTargetKindField = direction === "out" ? "to_kind" : "from_kind";
-  const directBranch = compileRecursiveBranch({
+  const directBranch = compileRecursiveBranch2({
     joinField: directJoinField,
     targetField: directTargetField,
     joinKindField: directJoinKindField,
@@ -3107,7 +3104,7 @@ function compileRecursiveCte(ast, traversal, graphId, ctx, requiredColumnsByAlia
       (kind) => directEdgeKinds.includes(kind)
     );
     const duplicateGuard = overlappingKinds.length > 0 ? drizzleOrm.sql`NOT (e.from_id = e.to_id AND ${compileKindFilter2(overlappingKinds, "e.kind")})` : void 0;
-    const inverseBranch = compileRecursiveBranch({
+    const inverseBranch = compileRecursiveBranch2({
       joinField: inverseJoinField,
       targetField: inverseTargetField,
       joinKindField: inverseJoinKindField,
@@ -3650,19 +3647,17 @@ function compilePrefixedSelect(ast, prefix, dialect) {
 function compileFieldValueForSetOp(source, prefix, alias, dialect) {
   if ("__type" in source && source.__type === "aggregate") {
     const { field: field2, function: fn } = source;
-    const cteName = `cte_${prefix}_${field2.alias}`;
     switch (fn) {
-      case "count": {
-        return drizzleOrm.sql`COUNT(${drizzleOrm.sql.raw(cteName)}.${drizzleOrm.sql.raw(field2.alias)}_id)`;
-      }
-      case "countDistinct": {
-        return drizzleOrm.sql`COUNT(DISTINCT ${drizzleOrm.sql.raw(cteName)}.${drizzleOrm.sql.raw(field2.alias)}_id)`;
-      }
+      case "count":
+      case "countDistinct":
       case "sum":
       case "avg":
       case "min":
       case "max": {
         const column = compileFieldColumnForSetOp(field2, prefix, dialect);
+        if (fn === "countDistinct") {
+          return drizzleOrm.sql`COUNT(DISTINCT ${column})`;
+        }
         return drizzleOrm.sql`${drizzleOrm.sql.raw(fn.toUpperCase())}(${column})`;
       }
       default: {
@@ -3876,9 +3871,7 @@ function collectRequiredColumnsByAlias2(ast) {
         markFieldRefAsRequired(requiredColumnsByAlias, source);
       } else {
         addRequiredColumn(requiredColumnsByAlias, source.field.alias, "id");
-        if (source.function !== "count" && source.function !== "countDistinct") {
-          markFieldRefAsRequired(requiredColumnsByAlias, source.field);
-        }
+        markFieldRefAsRequired(requiredColumnsByAlias, source.field);
       }
     }
   }
@@ -11195,6 +11188,142 @@ async function executeNodeBulkGetOrCreateByConstraint(ctx, kind, constraintName,
   }
   return results;
 }
+var DEFAULT_SUBGRAPH_MAX_DEPTH = 10;
+function normalizeProps(value) {
+  return typeof value === "string" ? value : JSON.stringify(value ?? {});
+}
+async function executeSubgraph(params) {
+  const { options } = params;
+  if (options.edges.length === 0) {
+    return { nodes: [], edges: [] };
+  }
+  const maxDepth = Math.min(
+    options.maxDepth ?? DEFAULT_SUBGRAPH_MAX_DEPTH,
+    MAX_RECURSIVE_DEPTH
+  );
+  const ctx = {
+    graphId: params.graphId,
+    rootId: params.rootId,
+    edgeKinds: options.edges,
+    maxDepth,
+    includeKinds: options.includeKinds,
+    excludeRoot: options.excludeRoot ?? false,
+    direction: options.direction ?? "out",
+    cyclePolicy: options.cyclePolicy ?? "prevent",
+    dialect: params.dialect,
+    schema: params.schema ?? DEFAULT_SQL_SCHEMA,
+    backend: params.backend
+  };
+  const reachableCte = buildReachableCte(ctx);
+  const includedIdsCte = buildIncludedIdsCte(ctx);
+  const [nodeRows, edgeRows] = await Promise.all([
+    fetchSubgraphNodes(ctx, reachableCte, includedIdsCte),
+    fetchSubgraphEdges(ctx, reachableCte, includedIdsCte)
+  ]);
+  const nodes = nodeRows.map(
+    (row) => rowToNode({ ...row, props: normalizeProps(row.props) })
+  );
+  const edges = edgeRows.map(
+    (row) => rowToEdge({ ...row, props: normalizeProps(row.props) })
+  );
+  return {
+    nodes,
+    edges
+  };
+}
+function buildReachableCte(ctx) {
+  const shouldTrackPath = ctx.cyclePolicy === "prevent";
+  const edgeKindFilter = compileKindFilter(drizzleOrm.sql.raw("e.kind"), ctx.edgeKinds);
+  const initialPath = shouldTrackPath ? ctx.dialect.initializePath(drizzleOrm.sql.raw("n.id")) : void 0;
+  const pathExtension = shouldTrackPath ? ctx.dialect.extendPath(drizzleOrm.sql.raw("r.path"), drizzleOrm.sql.raw("n.id")) : void 0;
+  const cycleCheck = shouldTrackPath ? ctx.dialect.cycleCheck(drizzleOrm.sql.raw("n.id"), drizzleOrm.sql.raw("r.path")) : void 0;
+  const baseColumns = [drizzleOrm.sql`n.id`, drizzleOrm.sql`n.kind`, drizzleOrm.sql`0 AS depth`];
+  if (initialPath !== void 0) {
+    baseColumns.push(drizzleOrm.sql`${initialPath} AS path`);
+  }
+  const baseCase = drizzleOrm.sql`SELECT ${drizzleOrm.sql.join(baseColumns, drizzleOrm.sql`, `)} FROM ${ctx.schema.nodesTable} n WHERE n.graph_id = ${ctx.graphId} AND n.id = ${ctx.rootId} AND n.deleted_at IS NULL`;
+  const recursiveColumns = [
+    drizzleOrm.sql`n.id`,
+    drizzleOrm.sql`n.kind`,
+    drizzleOrm.sql`r.depth + 1 AS depth`
+  ];
+  if (pathExtension !== void 0) {
+    recursiveColumns.push(drizzleOrm.sql`${pathExtension} AS path`);
+  }
+  const recursiveWhereClauses = [
+    drizzleOrm.sql`e.graph_id = ${ctx.graphId}`,
+    edgeKindFilter,
+    drizzleOrm.sql`e.deleted_at IS NULL`,
+    drizzleOrm.sql`n.deleted_at IS NULL`,
+    drizzleOrm.sql`r.depth < ${ctx.maxDepth}`
+  ];
+  if (cycleCheck !== void 0) {
+    recursiveWhereClauses.push(cycleCheck);
+  }
+  const forceWorktableOuterJoinOrder = ctx.dialect.capabilities.forceRecursiveWorktableOuterJoinOrder;
+  const recursiveCase = ctx.direction === "both" ? compileBidirectionalBranch({
+    recursiveColumns,
+    whereClauses: recursiveWhereClauses,
+    forceWorktableOuterJoinOrder,
+    schema: ctx.schema
+  }) : compileRecursiveBranch({
+    recursiveColumns,
+    whereClauses: recursiveWhereClauses,
+    joinField: "from_id",
+    targetField: "to_id",
+    targetKindField: "to_kind",
+    forceWorktableOuterJoinOrder,
+    schema: ctx.schema
+  });
+  return drizzleOrm.sql`WITH RECURSIVE reachable AS (${baseCase} UNION ALL ${recursiveCase})`;
+}
+function compileRecursiveBranch(params) {
+  const columns = [...params.recursiveColumns];
+  const selectClause = drizzleOrm.sql`SELECT ${drizzleOrm.sql.join(columns, drizzleOrm.sql`, `)}`;
+  const nodeJoin = drizzleOrm.sql`JOIN ${params.schema.nodesTable} n ON n.graph_id = e.graph_id AND n.id = e.${drizzleOrm.sql.raw(params.targetField)} AND n.kind = e.${drizzleOrm.sql.raw(params.targetKindField)}`;
+  if (params.forceWorktableOuterJoinOrder) {
+    const allWhere = [
+      ...params.whereClauses,
+      drizzleOrm.sql`e.${drizzleOrm.sql.raw(params.joinField)} = r.id`
+    ];
+    return drizzleOrm.sql`${selectClause} FROM reachable r CROSS JOIN ${params.schema.edgesTable} e ${nodeJoin} WHERE ${drizzleOrm.sql.join(allWhere, drizzleOrm.sql` AND `)}`;
+  }
+  const where = [...params.whereClauses];
+  return drizzleOrm.sql`${selectClause} FROM reachable r JOIN ${params.schema.edgesTable} e ON e.${drizzleOrm.sql.raw(params.joinField)} = r.id ${nodeJoin} WHERE ${drizzleOrm.sql.join(where, drizzleOrm.sql` AND `)}`;
+}
+function compileBidirectionalBranch(params) {
+  const columns = [...params.recursiveColumns];
+  const selectClause = drizzleOrm.sql`SELECT ${drizzleOrm.sql.join(columns, drizzleOrm.sql`, `)}`;
+  const nodeJoin = drizzleOrm.sql`JOIN ${params.schema.nodesTable} n ON n.graph_id = e.graph_id AND ((e.to_id = r.id AND n.id = e.from_id AND n.kind = e.from_kind) OR (e.from_id = r.id AND n.id = e.to_id AND n.kind = e.to_kind))`;
+  if (params.forceWorktableOuterJoinOrder) {
+    const allWhere = [
+      ...params.whereClauses,
+      drizzleOrm.sql`(e.from_id = r.id OR e.to_id = r.id)`
+    ];
+    return drizzleOrm.sql`${selectClause} FROM reachable r CROSS JOIN ${params.schema.edgesTable} e ${nodeJoin} WHERE ${drizzleOrm.sql.join(allWhere, drizzleOrm.sql` AND `)}`;
+  }
+  return drizzleOrm.sql`${selectClause} FROM reachable r JOIN ${params.schema.edgesTable} e ON (e.from_id = r.id OR e.to_id = r.id) ${nodeJoin} WHERE ${drizzleOrm.sql.join([...params.whereClauses], drizzleOrm.sql` AND `)}`;
+}
+function buildIncludedIdsCte(ctx) {
+  const filters = [];
+  if (ctx.includeKinds !== void 0 && ctx.includeKinds.length > 0) {
+    filters.push(compileKindFilter(drizzleOrm.sql.raw("kind"), ctx.includeKinds));
+  }
+  if (ctx.excludeRoot) {
+    filters.push(drizzleOrm.sql`id != ${ctx.rootId}`);
+  }
+  const whereClause = filters.length > 0 ? drizzleOrm.sql` WHERE ${drizzleOrm.sql.join(filters, drizzleOrm.sql` AND `)}` : drizzleOrm.sql``;
+  return drizzleOrm.sql`, included_ids AS (SELECT DISTINCT id FROM reachable${whereClause})`;
+}
+async function fetchSubgraphNodes(ctx, reachableCte, includedIdsCte) {
+  const query = drizzleOrm.sql`${reachableCte}${includedIdsCte} SELECT n.kind, n.id, n.props, n.version, n.valid_from, n.valid_to, n.created_at, n.updated_at, n.deleted_at FROM ${ctx.schema.nodesTable} n WHERE n.graph_id = ${ctx.graphId} AND n.id IN (SELECT id FROM included_ids)`;
+  return ctx.backend.execute(query);
+}
+async function fetchSubgraphEdges(ctx, reachableCte, includedIdsCte) {
+  const edgeKindFilter = compileKindFilter(drizzleOrm.sql.raw("e.kind"), ctx.edgeKinds);
+  const query = drizzleOrm.sql`${reachableCte}${includedIdsCte} SELECT e.id, e.kind, e.from_kind, e.from_id, e.to_kind, e.to_id, e.props, e.valid_from, e.valid_to, e.created_at, e.updated_at, e.deleted_at FROM ${ctx.schema.edgesTable} e WHERE e.graph_id = ${ctx.graphId} AND ${edgeKindFilter} AND e.deleted_at IS NULL AND e.from_id IN (SELECT id FROM included_ids) AND e.to_id IN (SELECT id FROM included_ids)`;
+  return ctx.backend.execute(query);
+}
 
 // src/store/store.ts
 var Store = class {
@@ -11398,6 +11527,39 @@ var Store = class {
   query() {
     return this.#createQueryForBackend(this.#backend);
   }
+  // === Subgraph Extraction ===
+  /**
+   * Extracts a typed subgraph by traversing from a root node.
+   *
+   * Performs a BFS traversal from `rootId` following the specified edge kinds,
+   * then returns all reachable nodes and the edges connecting them.
+   *
+   * @example
+   * ```typescript
+   * const result = await store.subgraph(run.id, {
+   *   edges: ["has_task", "runs_agent", "uses_skill"],
+   *   maxDepth: 4,
+   *   includeKinds: ["Run", "Task", "Agent", "Skill"],
+   * });
+   *
+   * for (const node of result.nodes) {
+   *   switch (node.kind) {
+   *     case "Task": console.log(node.name); break;
+   *     case "Agent": console.log(node.model); break;
+   *   }
+   * }
+   * ```
+   */
+  async subgraph(rootId, options) {
+    return executeSubgraph({
+      graphId: this.graphId,
+      rootId,
+      backend: this.#backend,
+      dialect: chunk2WVFEIHR_cjs.getDialect(this.#backend.dialect),
+      schema: this.#schema,
+      options
+    });
+  }
   // === Transactions ===
   /**
    * Executes a function within a transaction.