Merge pull request #33 from bruinenxyz/tyler/add-relate-backend

Add backend for relate, order, and take steps

backend/src/definitions/pipeline.ts

@@ -268,6 +268,7 @@ export type TakeStep = z.infer<typeof TakeStepSchema>;
 export const RelateStepSchema = z.object({
   type: z.literal(StepIdentifierEnum.Relate),
   relation: InferredSchemaRelationSchema,
+  inputSchema: z.array(InferredSchemaColumnSchema),
 });
 export type RelateStep = z.infer<typeof RelateStepSchema>;
 

backend/src/relations/relations.module.ts

@@ -8,5 +8,6 @@ import { TablesModule } from "@/tables/tables.module";
   controllers: [RelationsController],
   providers: [RelationsService],
   imports: [PostgresAdapterModule, TablesModule],
+  exports: [RelationsService],
 })
 export class RelationsModule {}

backend/src/user-queries/query-parsing/parse-aggregate.ts

@@ -33,6 +33,14 @@ export function parseAggregate(
   );
   assert(aggregateTable, `Table not found: ${aggregateStep.column.table}`);
 
+  // Prefix the aggregate column with the relation name if it exists
+  let aggregatePrefix;
+  if (aggregateStep.column.relation) {
+    aggregatePrefix = aggregateStep.column.relation.as;
+  } else {
+    aggregatePrefix = aggregateTable.external_name;
+  }
+
   let aggregatePrql = "";
   if (aggregateStep.group.length === 0) {
     // Aggregate over the whole table
@@ -41,12 +49,22 @@ export function parseAggregate(
       from: `step_${index - 1}`,
       as: aggregateStep.as,
       operation: aggregateStep.operation,
-      // TODO this might need to include the relation path in the name to avoid collisions
-      column: `${aggregateTable.external_name}__${aggregateStep.column.name}`,
+      column: `${aggregatePrefix}__${aggregateStep.column.name}`,
     });
   } else {
     // Aggregate over a group
     const groupedColumnNames = _.map(aggregateStep.group, (column) => {
+      // If the column was created via aggregate, we don't need to prefix it
+      if (column.table === "aggregate") {
+        return column.name;
+      }
+
+      // If the column was added via relation, prefix it with the relation name
+      if (column.relation) {
+        return `${column.relation.as}__${column.name}`;
+      }
+
+      // Otherwise, prefix it with its base table name
       const table = _.find(tables, (table) => table.id === column.table);
       assert(table, `Table not found: ${column.table}`);
 
@@ -59,8 +77,7 @@ export function parseAggregate(
       group: groupedColumnNames.join(", "),
       as: aggregateStep.as,
       operation: aggregateStep.operation,
-      // TODO this might need to include the relation path in the name to avoid collisions
-      column: `${aggregateTable.external_name}__${aggregateStep.column.name}`,
+      column: `${aggregatePrefix}__${aggregateStep.column.name}`,
     });
   }
 

backend/src/user-queries/query-parsing/parse-derive.ts

@@ -0,0 +1 @@
+// TODO derive step is not yet supported

backend/src/user-queries/query-parsing/parse-filter.ts

@@ -0,0 +1 @@
+// TODO filter step is not yet supported

backend/src/user-queries/query-parsing/parse-order.ts

@@ -0,0 +1,46 @@
+import { OrderStep, Table } from "@/definitions";
+import { compileTemplate } from "./utils";
+import * as _ from "lodash";
+import * as assert from "assert";
+
+const orderTemplate = `
+let {{ stepName }} = (
+  from {{ from }}
+  sort { {{ order }} }
+)`;
+
+export function parseOrder(
+  orderStep: OrderStep,
+  index: number,
+  tables: Table[],
+): string {
+  const orderPrql = compileTemplate(orderTemplate, {
+    stepName: `step_${index}`,
+    from: `step_${index - 1}`,
+    order: _.map(orderStep.order, (singleOrder) => {
+      const dir = singleOrder.direction === "asc" ? "+" : "-";
+
+      // If the ordered column was created via aggregate, we don't need to prefix it
+      if (singleOrder.column.table === "aggregate") {
+        return `${dir}${singleOrder.column.name}`;
+      }
+
+      // If the ordered column was added via relation, prefix it with the relation name
+      if (singleOrder.column.relation) {
+        return `${dir}${singleOrder.column.relation.as}__${singleOrder.column.name}`;
+      }
+
+      // Otherwise, prefix it with its base table name
+      const table = _.find(
+        tables,
+        (table) => table.id === singleOrder.column.table,
+      );
+      assert(table, `Table not found: ${singleOrder.column.table}`);
+
+      // TODO include schema?
+      return `${dir}${table.external_name}__${singleOrder.column.name}`;
+    }).join(", "),
+  });
+
+  return orderPrql;
+}

backend/src/user-queries/query-parsing/parse-query.ts

@@ -1,18 +1,23 @@
 import {
   ExternalColumn,
   Pipeline,
+  Relation,
   StepIdentifierEnum,
   Table,
 } from "@/definitions";
 import { parseAggregate } from "./parse-aggregate";
 import { parseFrom } from "./parse-from";
+import { parseOrder } from "./parse-order";
+import { parseRelate } from "./parse-relate";
 import { parseSelect } from "./parse-select";
+import { parseTake } from "./parse-take";
 import { compile } from "prql-js";
 import * as _ from "lodash";
 
 function objectToPrql(
   pipeline: Pipeline,
   tables: Table[],
+  relations: Relation[],
   tablesSchema: Record<string, Record<string, ExternalColumn>>,
 ): string {
   let prql = "prql target:sql.postgres\n";
@@ -23,9 +28,17 @@ function objectToPrql(
       case StepIdentifierEnum.Aggregate:
         prql += parseAggregate(step, index + 1, tables);
         break;
+      case StepIdentifierEnum.Order:
+        prql += parseOrder(step, index + 1, tables);
+        break;
+      case StepIdentifierEnum.Relate:
+        prql += parseRelate(step, index + 1, tables, relations, tablesSchema);
+        break;
       case StepIdentifierEnum.Select:
         prql += parseSelect(step, index + 1, tables);
         break;
+      case StepIdentifierEnum.Take:
+        prql += parseTake(step, index + 1);
       default:
         break;
     }
@@ -39,9 +52,10 @@ function objectToPrql(
 export function writeSQL(
   pipeline: Pipeline,
   tables: Table[],
+  relations: Relation[],
   tablesSchema: Record<string, Record<string, ExternalColumn>>,
 ) {
-  const prqlQuery = objectToPrql(pipeline, tables, tablesSchema);
+  const prqlQuery = objectToPrql(pipeline, tables, relations, tablesSchema);
 
   console.log("prqlQuery:");
   console.log(prqlQuery);

backend/src/user-queries/query-parsing/parse-relate.ts

@@ -0,0 +1,196 @@
+import {
+  ExternalColumn,
+  RelateStep,
+  Relation,
+  RelationTypeEnum,
+  Table,
+} from "@/definitions";
+import { baseObjectTemplate, compileTemplate } from "./utils";
+import * as _ from "lodash";
+import * as assert from "assert";
+
+const relateTemplate = `
+let {{ stepName }} = (
+  from {{ from }}
+  join side:left {{ relation }} (this.{{ baseJoinColumn }}=={{ relation }}.{{ relationJoinColumn }})
+)`;
+
+const relateTemplateManyToMany = `
+let {{ stepName }} = (
+  from {{ from }}
+  join side:left {{ joinTable }} (this.{{ baseJoinColumn }}=={{ joinTable }}.{{ joinTableBaseColumn }})
+  join side:left {{ relation }} (this.{{ joinTableRelationColumn }}=={{ relation }}.{{ relationJoinColumn }})
+  select { {{ selectProperties }} }
+)`;
+
+export function parseRelate(
+  relateStep: RelateStep,
+  index: number,
+  tables: Table[],
+  relations: Relation[],
+  tablesSchema: Record<string, Record<string, ExternalColumn>>,
+): string {
+  const relateStepPrql: string[] = [];
+
+  const relatedTable = _.find(
+    tables,
+    (table) => table.id === relateStep.relation.table,
+  );
+  assert(relatedTable, `Table not found: ${relateStep.relation.table}`);
+
+  const columns = _.keys(tablesSchema[relatedTable.external_name]);
+
+  // Add the base object variable to the PRQL, prefixing every column with the alias
+  // TODO we will need to add the schema to the names below
+  relateStepPrql.push(
+    compileTemplate(baseObjectTemplate, {
+      varName: relateStep.relation.as,
+      tableName: relatedTable.external_name,
+      deriveProperties: _.map(columns, (column) => {
+        return `${relateStep.relation.as}__${column} = ${column}`;
+      }).join(", "),
+      selectProperties: _.map(columns, (column) => {
+        return `${relateStep.relation.as}__${column}`;
+      }).join(", "),
+    }),
+  );
+
+  const relation = _.find(
+    relations,
+    (relation) => relation.id === relateStep.relation.relation,
+  );
+  assert(relation, `Relation not found: ${relateStep.relation.relation}`);
+
+  if (relation.type === RelationTypeEnum.ManyToMany) {
+    let baseTableId;
+    let baseColumnName;
+    let relationColumnName;
+    let joinTableBaseColumnName;
+    let joinTableRelationColumnName;
+
+    // Check which side of the relation is the base table and use the appropriate columns
+    if (relateStep.relation.table === relation.table_1) {
+      baseTableId = relation.table_2;
+      baseColumnName = relation.column_2;
+      relationColumnName = relation.column_1;
+      joinTableBaseColumnName = relation.join_column_2;
+      joinTableRelationColumnName = relation.join_column_1;
+    } else {
+      baseTableId = relation.table_1;
+      baseColumnName = relation.column_1;
+      relationColumnName = relation.column_2;
+      joinTableBaseColumnName = relation.join_column_1;
+      joinTableRelationColumnName = relation.join_column_2;
+    }
+
+    // Pull the external name of the join table
+    const joinTable = _.find(
+      tables,
+      (table) => table.id === relation.join_table,
+    );
+    assert(joinTable, `Table not found: ${relation.join_table}`);
+
+    const joinTableColumns = _.keys(tablesSchema[joinTable.external_name]);
+
+    relateStepPrql.push(
+      compileTemplate(baseObjectTemplate, {
+        // TODO add schema
+        varName: `${relateStep.relation.as}_join_table`,
+        tableName: joinTable.external_name,
+        deriveProperties: _.map(joinTableColumns, (column) => {
+          return `${relateStep.relation.as}_join_table__${column} = ${column}`;
+        }).join(", "),
+        selectProperties: _.map(joinTableColumns, (column) => {
+          return `${relateStep.relation.as}_join_table__${column}`;
+        }).join(", "),
+      }),
+    );
+
+    // Check if the base join key is on a table that has been related in previously
+    let baseColumnPrefix;
+    if (relateStep.relation.on.relation) {
+      baseColumnPrefix = relateStep.relation.on.relation.as;
+    } else {
+      const baseTable = _.find(tables, (table) => table.id === baseTableId);
+      assert(baseTable, `Table not found: ${baseTableId}`);
+      baseColumnPrefix = baseTable.external_name;
+    }
+
+    // Select all of the properties on the base and related tables (excluding joiun table columns)
+    // TODO add schema
+    const selectProperties = _.concat(
+      _.map(relateStep.inputSchema, (column) => {
+        if (column.table === "aggregate") {
+          return column.name;
+        }
+
+        if (column.relation) {
+          return `${column.relation.as}__${column.name}`;
+        }
+
+        const table = _.find(tables, (table) => table.id === column.table);
+        assert(table, `Table not found: ${column.table}`);
+
+        return `${table.external_name}__${column.name}`;
+      }),
+      _.map(columns, (column) => {
+        return `${relateStep.relation.as}__${column}`;
+      }),
+    ).join(", ");
+
+    // Add the relation to the step PRQL
+    relateStepPrql.push(
+      compileTemplate(relateTemplateManyToMany, {
+        stepName: `step_${index}`,
+        from: `step_${index - 1}`,
+        joinTable: `${relateStep.relation.as}_join_table`,
+        baseJoinColumn: `${baseColumnPrefix}__${baseColumnName}`,
+        joinTableBaseColumn: `${relateStep.relation.as}_join_table__${joinTableBaseColumnName}`,
+        joinTableRelationColumn: `${relateStep.relation.as}_join_table__${joinTableRelationColumnName}`,
+        relation: relateStep.relation.as,
+        relationJoinColumn: `${relateStep.relation.as}__${relationColumnName}`,
+        // TODO update
+        selectProperties: selectProperties,
+      }),
+    );
+  } else {
+    let baseTableId;
+    let baseColumnName;
+    let relationColumnName;
+
+    // Check which side of the relation is the base table and use the appropriate columns
+    if (relateStep.relation.table === relation.table_1) {
+      baseTableId = relation.table_2;
+      baseColumnName = relation.column_2;
+      relationColumnName = relation.column_1;
+    } else {
+      baseTableId = relation.table_1;
+      baseColumnName = relation.column_1;
+      relationColumnName = relation.column_2;
+    }
+
+    // Check if the base join key is on a table that has been related in previously
+    let baseColumnPrefix;
+    if (relateStep.relation.on.relation) {
+      baseColumnPrefix = relateStep.relation.on.relation.as;
+    } else {
+      const baseTable = _.find(tables, (table) => table.id === baseTableId);
+      assert(baseTable, `Table not found: ${baseTableId}`);
+      baseColumnPrefix = baseTable.external_name;
+    }
+
+    // Add the relation to the step PRQL
+    // TODO add in schema to the naming
+    relateStepPrql.push(
+      compileTemplate(relateTemplate, {
+        stepName: `step_${index}`,
+        from: `step_${index - 1}`,
+        relation: relateStep.relation.as,
+        baseJoinColumn: `${baseColumnPrefix}__${baseColumnName}`,
+        relationJoinColumn: `${relateStep.relation.as}__${relationColumnName}`,
+      }),
+    );
+  }
+
+  return relateStepPrql.join("\n");
+}

backend/src/user-queries/query-parsing/parse-select.ts

@@ -21,10 +21,17 @@ export function parseSelect(
     stepName: `step_${index}`,
     from: `step_${index - 1}`,
     columns: _.map(selectStep.select, (column) => {
+      // If the column was created via aggregate, we don't need to prefix it
       if (column.table === "aggregate") {
         return column.name;
       }
 
+      // If the column was added via relation, prefix it with the relation name
+      if (column.relation) {
+        return `${column.relation.as}__${column.name}`;
+      }
+
+      // Otherwise, prefix it with its base table name
       const table = _.find(tables, (table) => table.id === column.table);
       assert(table, `Table not found: ${column.table}`);