Zenstack is using base Prisma queries instead of extended ones, when overriding queries during Prisma.Extensions

[CollinKempkes]

Description and expected behavior
We are trying to integrate some internal mapping of currencies. Therefore we committed ourselves to a currency length of five, if some currencies are shorter they will be prepended by Z until they reach that amount of chars. We do this by overriding mutation queries inside prisma extensions. Additionally we want to use inside of the app the normal currencies (so we don't have to map all these custom currencies in the runtime, when we do API calls to service providers). So we have some kind of pre- and postTransformations of queries.

currency-extension.ts

import { Prisma } from '@prisma/client';
import {
  postTransformations,
  preTransformations,
} from '@vdrip-database/database.util-test';

export const CurrencyTransformationExtensionTest = Prisma.defineExtension({
  name: 'Currency Transformation',

  query: {
    $allModels: {
      async $allOperations({ operation, args, query }) {
        const valueFields = ['data', 'create', 'update', 'where'];

        switch (operation) {
          // For all select operations
          case 'aggregate':
          case 'count':
          case 'findFirst':
          case 'findFirstOrThrow':
          case 'findMany':
          case 'findUnique':
          case 'groupBy':
          case 'upsert':
          case 'update':
          case 'updateMany':
          case 'findUniqueOrThrow':
          // For all mutation operations
          case 'create':
          case 'createMany':
          case 'update':
          case 'updateMany':
          case 'upsert': {
            valueFields.forEach((field) => {
              // @ts-expect-error
              if (args[field]) {
                // @ts-expect-error
                args[field] = preTransformations.args.currencyTransform(
                  // @ts-expect-error
                  args[field]
                );
              }
            });
          }
        }

        return postTransformations.result.currencyTransform(await query(args));
      },
    },
  },
});

Override of the currency field is happening here:
preTransformations.args.currencyTransform(...)

Here is the impl of that:
database.util.ts

import { Prisma } from '@prisma/client';
import {
  convertCurrencyToVdripCurrency as convertCurrencyToInternalCurrency,
  convertVdripCurrencyToCurrency,
} from '@vdrip-app/stripe/payment.utils';
import { containsAttribute } from '@vdrip-app/utils/objects';

export const doesModelIncludeField = (model: string, fieldName: string) => {
  const dbModels = Prisma.dmmf.datamodel.models;
  const dbModel = dbModels.find((dbModel) => dbModel.name === model);
  return dbModel?.fields.some((field) => field.name === fieldName);
};

export const preTransformations = {
  args: {
    currencyTransform: (data: any) => {
      const currencyFields = ['currency'];

      currencyFields.forEach((field) => {
        if (data && data[field]) {
          data[field] = convertCurrencyToInternalCurrency(data[field]);
        }
      });

      return data;
    },
    softDeleteTransform: (model: string, column: string, where: any) => {
      const isDeleteable = doesModelIncludeField(model, column);

      if (isDeleteable && !containsAttribute(column, where)) {
        where = { ...where, deleted_at: null };
      }

      return where;
    },
  },
};

export const postTransformations = {
  result: {
    currencyTransform: (result: any) => {
      if (!result) {
        return result;
      }

      const currencyFields = ['currency'];

      currencyFields.forEach((field) => {
        if (Array.isArray(result)) {
          result = result.map((result) => {
            if (result[field]) {
              result[field] = convertVdripCurrencyToCurrency(result[field]);
            }
            return result;
          });
        }

        if (result[field]) {
          result[field] = convertVdripCurrencyToCurrency(result[field]);
        }
      });

      return result;
    },
  },
};

Additionally we override the response got by prisma return postTransformations.result.currencyTransform(await query(args));. The problem seems to be that an enhanced + extended prisma client is using base methods of Prisma.
Therefore this test scenario:
test.spec.ts

it('should transform currencies automatically, but save them differently', async () => {
        // setup
        const extendedAndEnhancedPrisma = enhance(testContext.prismaClient);
        const unextendedEnhancedPrisma = enhance(unextendedPrismaClient);

        const testEntity = await testContext.prismaClient.test.create({
          data: {
            currency: 'USD',
          },
        });
        testEntity;
        //! ##### The Code is working until here #####

        const enhancedTestEntity = await extendedAndEnhancedPrisma.test.create({
          data: {
            currency: 'USD',
          },
        });

        // effects
        const transformedEntity =
          await extendedAndEnhancedPrisma.test.findFirst({
            where: {
              id: enhancedTestEntity.id,
            },
          });
        const realEntity = await unextendedEnhancedPrisma.test.findFirst({
          where: {
            id: enhancedTestEntity.id,
          },
        });

        // post-conditions
        expect(transformedEntity).toEqual(
          expect.objectContaining({
            currency: enhancedTestEntity.currency,
          })
        );
        expect(realEntity).toEqual(
          expect.objectContaining({
            currency: convertCurrencyToVdripCurrency(
              enhancedTestEntity.currency as AcceptedPresentmentCurrencies
            ),
          })
        );
      });

Everything really looks fine until the enhanced create. the preTransformation is working (in the db we have ZZUSD) saved and the postTransformation is working (when looking at testEntity it is saying that currency is USD).

For completeness, here the zmodel:
test.zmodel

model Test {
  id String @id @default(dbgenerated("uuid_generate_v4()")) @db.Uuid
  currency String

  @@validate(currency in [
      'ZZEUR',
      'ZZUSD'
    ],
    'type must be one of the following values: ZZEUR, ZZUSD'
  )

  @@schema("public")
  @@map("test")

  @@allow("all", true)
}

The error that is occurring inside extendedAndEnhancedPrisma.test.create(...) is

Error calling enhanced Prisma method `create`: denied by policy: test entities failed 'create' check, input failed validation: Validation error: type must be one of the following values: ZZEUR, ZZUSD

Environment:

• ZenStack version: 1.11.1
• Prisma version: 5.9.1
• Database type: Postgresql

Additional context
Dit not try yet to use new model functions instead of query overrides, will try this now and will comment the outcome.
If you need more code example just hit me up :)
The transformer of the currency filling it pretty basic for now we can just say it adds static 'ZZ' to the currency when writing and does .slice(2), when reading.

[CollinKempkes]

Addition:
Seems to work when creating with completely custom model functions:
test.spec.ts

...
const testEntity =
          await testContext.prismaClient.test.createWithCurrency({
            data: {
              currency: 'USD',
            },
          });
        testEntity;
        const enhancedTestEntity =
          await extendedAndEnhancedPrisma.test.createWithCurrency({
            data: {
              currency: 'USD',
            },
          });
...

currency-extension.ts

...
model: {
    $allModels: {
      async createWithCurrency<Model, Args extends object>(
        this: Model,
        args: Prisma.Exact<Args, Prisma.Args<Model, 'create'>>
      ): Promise<Prisma.Result<Model, Args, 'create'>> {
        // @ts-expect-error
        args.data = preTransformations.args.currencyTransform(
          // @ts-expect-error
          args.data
        );
        // @ts-expect-error: Requires more types from Prisma
        return this.create({
          ...args,
        });
      },
    },
  },
...

This test is running green now. But I would say the desired behavior is that it also works with query overrides instead of model query creations.

[ymc9]

Hi @CollinKempkes , thank you for filing the issue with great details!

Could you try to enhance the client first and then install the extension? The reason is that both Prisma client extensions and ZenStack enhancements are Javascript proxies around the original prisma client. In your current setup, the client extension proxy directly wraps around the client, and then ZenStack wraps outside of it. So when you make a create call, ZenStack's proxy receives it first, and the args haven't been transformed yet. The raw args are then rejected by the validation rules.

Extending with a new createWithCurrency method works because that method is not intercepted by ZenStack, so it reaches the client extension directly.

I understand it may be a bit confusing. Let me know if you feel more clarification is needed.

[CollinKempkes]

Thanks for your response @ymc9, the explanation really makes sense, thx for it!
But when I reorder the sequence of extending and enhancing I face another issue. For some reason my validations and permissions are skipped then (Maybe they will get overriden by the extension?).

Basically I ran this test to check if everything works fine:

it('test', async () => {
    // setup
    const extendedUser =
      await testContext.prismaClient.user.findFirstPermissionsAndRestrictions({
        where: {
          id: createdEntities.sellerUser.id,
        },
      });
    const enhancedExtendedPrisma = new ExtendedPrismaService(
      enhance(new PrismaClient(), {
        user: extendedUser,
      })
    ).withExtensions();
    const extendedEnhancedPrisma = enhance(
      new ExtendedPrismaService(new PrismaClient()).withExtensions(),
      {
        user: extendedUser,
      }
    );

    const test = await enhancedExtendedPrisma.test.create({
      data: {
        currency: 'DDDDD',
      },
    });
    test;

    // ### The following query will break ###
    const test2 = await extendedEnhancedPrisma.test.create({
      data: {
        currency: 'DDDDD',
      },
    });
    test2;
  });

model Test {
  id String @id @default(dbgenerated("uuid_generate_v4()")) @db.Uuid
  currency String

  @@validate(currency in [
      'ZZEUR',
      'ZZUSD'
    ],
    'type must be one of the following values: ZZEUR, ZZUSD'
  )

  @@schema("public")
  @@map("test")

  @@allow("all", true)
}

the query for test2 will break with the following reason, which is totally correct: Error calling enhanced Prisma method 'create': denied by policy: test entities failed 'create' check, input failed validation: Validation error: type must be one of the following values: ZZEUR, ZZUSD
But test is running through, which should not be the case as the validation should break.

same thing for permissions:

model Test {
  id String @id @default(dbgenerated("uuid_generate_v4()")) @db.Uuid
  currency String

  @@schema("public")
  @@map("test")

  @@allow("read", true)
  @@allow("create", currency == 'ZZEUR' || currency == 'ZZUSD')
}

Now the second query will fail with the error: Error calling enhanced Prisma method 'create': denied by policy: test entities failed 'create' check

Just to ensure that all information are given, here is the code of the extended-prisma-service:

Expand me for ExtendedPrismaService

...
export class ExtendedPrismaService {
  constructor(
    @Inject(PrismaClient) private readonly prismaClient: PrismaClient
  ) {}

  withExtensions() {
    return this.prismaClient
      .$extends({
        name: 'Extended Prisma Client',
        model: {
          $allModels: {
            findFirstNotNull,
          },
          identity: {
            findFirstWithRevokableListingsAndAsks,
          },
          localizedProperty: {
            findFirstWithAllRelations,
          },
          user: {
            findFirstPermissionsAndRestrictions,
          },
        },
      })
      .$extends(SoftDeleteExtension)
      .$extends(CurrencyTransformationExtension);
  }
}

[ymc9]

It seems when query is called inside an extension, Prisma somehow lets the original client handle the call (instead of the client instance with which .$extends is called) ... I'll need to debug deeper to understand why.

[ymc9]

Hi @CollinKempkes , I debugged through some of Prisma's code and felt there isn't a good way to get it to work, when client extensions are installed onto an enhanced prisma client - at least for "query" extension. Internally Prisma calls into the original client and bypasses ZenStack's enhancement. My apologies for giving an incorrect suggestion initially.

I think a possible (a bit more complex) solution is to turn the client extension into a regular JS proxy. Something like:

const enhanced = enhance(prisma, undefined, { logPrismaQuery: true });

const extendedAndEnhancedPrisma = new Proxy(enhanced, {
    get(target, prop, receiver) {
        const allModels = ['test'];

        if (!allModels.includes(prop as string)) {
            // not a model field access
            return Reflect.get(target, prop, receiver);
        }

        const value = Reflect.get(target, prop, receiver);
        if (value && typeof value === 'object') {
            // proxy the model-level prisma client
            return new Proxy(value, {
                get(target, prop, receiver) {
                    const valueFields = ['data', 'create', 'update', 'where'];

                    switch (prop) {
                        // For all select operations
                        case 'aggregate':
                        case 'count':
                        case 'findFirst':
                        case 'findFirstOrThrow':
                        case 'findMany':
                        case 'findUnique':
                        case 'groupBy':
                        case 'upsert':
                        case 'update':
                        case 'updateMany':
                        case 'findUniqueOrThrow':
                        // For all mutation operations
                        case 'create':
                        case 'createMany':
                        case 'update':
                        case 'updateMany':
                        case 'upsert': {
                            return async (args: any) => {
                                valueFields.forEach((field) => {
                                    if (args[field]) {
                                        args[field] = preTransformations.args.currencyTransform(args[field]);
                                    }
                                });
                                const result = await Reflect.get(target, prop, receiver)(args);
                                return postTransformations.result.currencyTransform(result);
                            };
                        }
                    }

                    return Reflect.get(target, prop, receiver);
                },
            });
        }
    },
});

[CollinKempkes]

No worries and thanks for your new suggestion. It seems like it works, but it does not work with $transactions.
Will figure some solution out and will post it here, if I am able to solve it

[ymc9]

Thank you!

Btw, I've added a dedicated documentation for client extensions here: https://zenstack.dev/docs/guides/client-extensions

[CollinKempkes]

Thanks for adding the docs! :)

Just a short update, it seems like it is working with this code here:

import { Prisma } from '@prisma/client';
...
import { isNil } from 'ramda';

/**
 * `withCurrencyFormat()` extends all queries of the prisma client to apply
 * currency code tranformations in both persistence and retrieval.
 */
export const withCurrencyFormat = <Client extends object>(
  client: Client
): Client => {
  return new Proxy(client, {
    get(target, prop, receiver) {
      const reflected = Reflect.get(target, prop, receiver);
      const allowedMethods = [ // TODO: Check if there is a better way than explicitly mentioning the allowed methods to change
        'aggregate',
        'count',
        'findFirst',
        'findFirstOrThrow',
        'findMany',
        'findUnique',
        'groupBy',
        'upsert',
        'update',
        'updateMany',
        'findUniqueOrThrow',
        'create',
        'createMany',
        'update',
        'updateMany',
        'upsert',
      ];

      const isTransaction = prop === '$transaction';

      /**
       * Recursively proxy when we are running a transaction.
       */
      if (isTransaction) {
        return async (...args: any) => {
          const [callback] = args;

          return callback(withCurrencyFormat(client));
        };
      }

      const [c1, ...model] = prop as string;
      const isModel = Object.values(Prisma.ModelName).includes(
        `${c1?.toUpperCase()}${model.join('')}` as Prisma.ModelName
      );

      /**
       * Recursively proxy when we are running a model query.
       */
      if (isModel) return withCurrencyFormat(reflected as Client);
      if (!allowedMethods.includes(prop as string)) return reflected; // This is mandatory to not override Prisma internal calls

      /**
       * Proxy specific operations
       */
      if (typeof reflected === 'function') {
        const fields = new Set(['data', 'create', 'update', 'where']);

        return async (args: any) => {
          fields.forEach((field) => {
            if (args?.[field]) {
              args[field] = preTransformations.args.currencyTransform(
                args[field]
              );
            }
          });

          // @ts-expect-error
          const result = await reflected(args);
          if (isNil(result)) return result;

          return Array.isArray(result)
            ? result.map((data) =>
                postTransformations.result.currencyTransform(data)
              )
            : postTransformations.result.currencyTransform(result);
        };
      }

      return reflected;
    },
  });
};

The sequence ob building the client is like this:

1. Extend Prisma Client
2. Enhance Client
3. Add proxies to client (to override inputs/ results of queries)

On first tests it seemed to have problems with the newest version 1.12.0 of prisma, but with 1.11.1 it seems to work fine, will add a new comment if I know whats happening there.

I will close this issue after I resolved all other open todos. Just wanted to add this code here, so you know that there is a way to make it work :)

[CollinKempkes]

Alright it seems to work out now after adding some changes to the transaction adjustment.

import { Prisma } from '@prisma/client';
import { DMMF } from '@prisma/generator-helper';
...
import { isNil } from 'ramda';

/**
 * `withCurrencyFormat()` extends all queries of the prisma client to apply
 * currency code tranformations in both persistence and retrieval.
 */
export const withCurrencyFormat = <Client extends object>(
  client: Client
): Client => {
  return new Proxy(client, {
    get(target, prop, receiver) {
      const reflected: any = Reflect.get(target, prop, receiver);
      const baseModelQueries = Object.keys(DMMF.ModelAction);

      const isTransaction = prop === '$transaction';

      /**
       * Use proxy for transactions callback.
       */
      if (isTransaction) {
        return (callback: any, ...rest: any[]) => {
          if (typeof callback !== 'function') {
            throw new Error('A function value input is expected');
          }

          return reflected.bind(target)((tx: any) => {
            return callback(withCurrencyFormat(tx));
          }, ...rest);
        };
      }

      const [c1, ...model] = prop as string;
      const isModel = Object.values(Prisma.ModelName).includes(
        `${c1?.toUpperCase()}${model.join('')}` as Prisma.ModelName
      );

      /**
       * Recursively proxy when we are running a model query.
       */
      if (isModel) return withCurrencyFormat(reflected as Client);
      if (!baseModelQueries.includes(prop as Prisma.DMMF.ModelAction))
        return reflected;

      /**
       * Proxy specific operations
       */
      if (typeof reflected === 'function') {
        const fields = new Set(['data', 'create', 'update', 'where']);

        return async (args: any) => {
          fields.forEach((field) => {
            if (args?.[field]) {
              args[field] = preTransformations.args.currencyTransform(
                args[field]
              );
            }
          });

          const result = await reflected(args);
          if (isNil(result)) return result;

          return postTransformations.result.currencyTransform(result);
        };
      }

      return reflected;
    },
  });
};

Handling arrays vs objects vs attributes is done inside the transformer scripts

[ymc9]

Thanks for sharing it. Great to know it's working now!