more robust alias handling

Author: ymc9

packages/runtime/src/client/query-utils.ts

@@ -56,7 +56,7 @@ export function requireField(schema: SchemaDef, modelOrType: string, field: stri
 
 export function getIdFields<Schema extends SchemaDef>(schema: SchemaDef, model: GetModels<Schema>) {
     const modelDef = requireModel(schema, model);
-    return modelDef?.idFields as GetModels<Schema>[];
+    return modelDef?.idFields;
 }
 
 export function requireIdFields(schema: SchemaDef, model: string) {

packages/runtime/src/plugins/policy/policy-handler.ts

@@ -32,7 +32,7 @@ import type { ClientContract } from '../../client';
 import type { CRUD } from '../../client/contract';
 import { getCrudDialect } from '../../client/crud/dialects';
 import type { BaseCrudDialect } from '../../client/crud/dialects/base-dialect';
-import { InternalError } from '../../client/errors';
+import { InternalError, QueryError } from '../../client/errors';
 import type { ProceedKyselyQueryFunction } from '../../client/plugin';
 import { getIdFields, requireField, requireModel } from '../../client/query-utils';
 import { ExpressionUtils, type BuiltinType, type Expression, type GetModels, type SchemaDef } from '../../schema';
@@ -73,7 +73,7 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         }
 
         let mutationRequiresTransaction = false;
-        const mutationModel = this.getMutationModel(node);
+        const { mutationModel } = this.getMutationModel(node);
 
         if (InsertQueryNode.is(node)) {
             // reject create if unconditional deny
@@ -168,7 +168,7 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         }
 
         // build a nested query with policy filter applied
-        const filter = this.buildPolicyFilter(table.model, undefined, 'read');
+        const filter = this.buildPolicyFilter(table.model, table.alias, 'read');
         const nestedSelect: SelectQueryNode = {
             kind: 'SelectQueryNode',
             from: FromNode.create([node.table]),
@@ -188,8 +188,8 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
 
         if (onConflict?.updates) {
             // for "on conflict do update", we need to apply policy filter to the "where" clause
-            const mutationModel = this.getMutationModel(node);
-            const filter = this.buildPolicyFilter(mutationModel, undefined, 'update');
+            const { mutationModel, alias } = this.getMutationModel(node);
+            const filter = this.buildPolicyFilter(mutationModel, alias, 'update');
             if (onConflict.updateWhere) {
                 onConflict = {
                     ...onConflict,
@@ -216,7 +216,8 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
             return result;
         } else {
             // only return ID fields, that's enough for reading back the inserted row
-            const idFields = getIdFields(this.client.$schema, this.getMutationModel(node));
+            const { mutationModel } = this.getMutationModel(node);
+            const idFields = getIdFields(this.client.$schema, mutationModel);
             return {
                 ...result,
                 returning: ReturningNode.create(
@@ -228,8 +229,8 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
 
     protected override transformUpdateQuery(node: UpdateQueryNode) {
         const result = super.transformUpdateQuery(node);
-        const mutationModel = this.getMutationModel(node);
-        let filter = this.buildPolicyFilter(mutationModel, undefined, 'update');
+        const { mutationModel, alias } = this.getMutationModel(node);
+        let filter = this.buildPolicyFilter(mutationModel, alias, 'update');
 
         if (node.from) {
             // for update with from (join), we need to merge join tables' policy filters to the "where" clause
@@ -247,8 +248,8 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
 
     protected override transformDeleteQuery(node: DeleteQueryNode) {
         const result = super.transformDeleteQuery(node);
-        const mutationModel = this.getMutationModel(node);
-        let filter = this.buildPolicyFilter(mutationModel, undefined, 'delete');
+        const { mutationModel, alias } = this.getMutationModel(node);
+        let filter = this.buildPolicyFilter(mutationModel, alias, 'delete');
 
         if (node.using) {
             // for delete with using (join), we need to merge join tables' policy filters to the "where" clause
@@ -272,19 +273,20 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         if (!node.returning) {
             return true;
         }
-        const idFields = getIdFields(this.client.$schema, this.getMutationModel(node));
+        const { mutationModel } = this.getMutationModel(node);
+        const idFields = getIdFields(this.client.$schema, mutationModel);
         const collector = new ColumnCollector();
         const selectedColumns = collector.collect(node.returning);
         return selectedColumns.every((c) => idFields.includes(c));
     }
 
     private async enforcePreCreatePolicy(node: InsertQueryNode, proceed: ProceedKyselyQueryFunction) {
-        const model = this.getMutationModel(node);
+        const { mutationModel } = this.getMutationModel(node);
         const fields = node.columns?.map((c) => c.column.name) ?? [];
-        const valueRows = node.values ? this.unwrapCreateValueRows(node.values, model, fields) : [[]];
+        const valueRows = node.values ? this.unwrapCreateValueRows(node.values, mutationModel, fields) : [[]];
         for (const values of valueRows) {
             await this.enforcePreCreatePolicyForOne(
-                model,
+                mutationModel,
                 fields,
                 values.map((v) => v.node),
                 proceed,
@@ -431,17 +433,13 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         }
 
         // do a select (with policy) in place of returning
-        const table = this.getMutationModel(node);
-        if (!table) {
-            throw new InternalError(`Unable to get table name for query node: ${node}`);
-        }
-
-        const idConditions = this.buildIdConditions(table, result.rows);
-        const policyFilter = this.buildPolicyFilter(table, undefined, 'read');
+        const { mutationModel } = this.getMutationModel(node);
+        const idConditions = this.buildIdConditions(mutationModel, result.rows);
+        const policyFilter = this.buildPolicyFilter(mutationModel, undefined, 'read');
 
         const select: SelectQueryNode = {
             kind: 'SelectQueryNode',
-            from: FromNode.create([TableNode.create(table)]),
+            from: FromNode.create([TableNode.create(mutationModel)]),
             where: WhereNode.create(conjunction(this.dialect, [idConditions, policyFilter])),
             selections: node.returning.selections,
         };
@@ -470,13 +468,23 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
 
     private getMutationModel(node: InsertQueryNode | UpdateQueryNode | DeleteQueryNode) {
         const r = match(node)
-            .when(InsertQueryNode.is, (node) => getTableName(node.into) as GetModels<Schema>)
-            .when(UpdateQueryNode.is, (node) => getTableName(node.table) as GetModels<Schema>)
+            .when(InsertQueryNode.is, (node) => ({
+                mutationModel: getTableName(node.into) as GetModels<Schema>,
+                alias: undefined,
+            }))
+            .when(UpdateQueryNode.is, (node) => {
+                if (!node.table) {
+                    throw new QueryError('Update query must have a table');
+                }
+                const r = this.extractTableName(node.table);
+                return r ? { mutationModel: r.model, alias: r.alias } : undefined;
+            })
             .when(DeleteQueryNode.is, (node) => {
                 if (node.from.froms.length !== 1) {
-                    throw new InternalError('Only one from table is supported for delete');
+                    throw new QueryError('Only one from table is supported for delete');
                 }
-                return getTableName(node.from.froms[0]) as GetModels<Schema>;
+                const r = this.extractTableName(node.from.froms[0]!);
+                return r ? { mutationModel: r.model, alias: r.alias } : undefined;
             })
             .exhaustive();
         if (!r) {
@@ -531,18 +539,18 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         return combinedPolicy;
     }
 
-    private extractTableName(from: OperationNode): { model: GetModels<Schema>; alias?: string } | undefined {
-        if (TableNode.is(from)) {
-            return { model: from.table.identifier.name as GetModels<Schema> };
+    private extractTableName(node: OperationNode): { model: GetModels<Schema>; alias?: string } | undefined {
+        if (TableNode.is(node)) {
+            return { model: node.table.identifier.name as GetModels<Schema> };
         }
-        if (AliasNode.is(from)) {
-            const inner = this.extractTableName(from.node);
+        if (AliasNode.is(node)) {
+            const inner = this.extractTableName(node.node);
             if (!inner) {
                 return undefined;
             }
             return {
                 model: inner.model,
-                alias: IdentifierNode.is(from.alias) ? from.alias.name : undefined,
+                alias: IdentifierNode.is(node.alias) ? node.alias.name : undefined,
             };
         } else {
             // this can happen for subqueries, which will be handled when nested