The open-source extension code is here, complete with a runnable example.

Introduction

Prisma, the popular Typescript ORM, has supported extending its client since version 4.1.0. The extension mechanism is straightforward and very flexible. You can use it to alter models, alter queries, add client level methods, and alter result data. Additionally, extensions can be chained so you can use multiple extensions together.

As you can see from a short list of available extensions here, there are myriad use cases for client extensions including data transformations, logging, caching, and additional security. You can even use it to support natural language queries via ChatGPT.

I had to bounce around a bit between their docs, examples, and discord to get our extension fully working, so I thought it would be worth walking through it.

Background

At Streamdal, we provide an open-source platform for building and executing Code-Native data pipelines. We already have a variety of SDKs to let our users view and transform data coming in and out of their Prisma client in the language of their choice. However, providing an extension allows our users to add any number of data pipelines to any existing Prisma operation simply by supplying a little extra configuration.

Getting Started

Prisma provides an extension starter template here. Fork or click the This Template button to create an extension project skeleton. By default, the starter template provides an example schema and implements a new existsFn on all Prisma models. You’ll find it at src/index.ts.

export const existsFn = (_extensionArgs: Args) =>
  Prisma.defineExtension({
    name: "prisma-extension-find-or-create",
    model: {
      $allModels: {
        async exists<T, A>(
          this: T,
          args: Prisma.Exact<A, Prisma.Args<T, 'findFirst'>>
        ): Promise<boolean> {

          const ctx = Prisma.getExtensionContext(this)
          const result = await (ctx as any).findFirst(args)
          return result !== null
        },
      },
    },
  })

In addition to extending models by adding new methods, you can add new client level methods or alter existing client operations such as find and create. We’ll be using this last technique as we want to allow our users to add pipelines to existing Prisma operations by adding a bit of configuration instrumentation and without having to rewrite a lot of code. Let’s try extending Prisma’s create. For a complete example see here.

import { Prisma } from "@prisma/client/extension";
import {
  Audience,
  OperationType,
  SDKResponse,
  Streamdal,
  StreamdalConfigs,
} from "@streamdal/node-sdk";

export { Audience, OperationType, StreamdalConfigs };

export type StreamdalArgs = {
  streamdalAudience?: Audience;
};

export const streamdal = (streamdalConfigs: StreamdalConfigs) => {
  const streamdal = new Streamdal(streamdalConfigs);
  const decoder = new TextDecoder();
  const encoder = new TextEncoder();

  return Prisma.defineExtension({
    name: "prisma-extension-streamdal",
    model: {
      $allModels: {
        async create<T, A>(
          this: T,
          args?: Prisma.Exact<A, Prisma.Args<T, "create">> & StreamdalArgs,
        ): Promise<Prisma.Result<T, A, "create">> {
          const { streamdalAudience, ...rest }: any & StreamdalArgs =
            args || {};
          const ctx = Prisma.getExtensionContext(this);

          if (streamdalAudience) {
            const streamdalResult: SDKResponse = await streamdal.process({
              audience: streamdalAudience,
              data: encoder.encode(JSON.stringify(rest.data)),
            });
            const data = decoder.decode(streamdalResult.data);
            return (ctx as any).$parent[ctx.$name as any].create({
              data: JSON.parse(data),
            });
          }

          return (ctx as any).$parent[ctx.$name as any].create(rest);
        },
      },
    },
  });
};

There’s quite a bit going on here so I’ll walk through it a bit at a time. I’ve included the imports from our Typescript Node SDK so you can see that adding libraries to a Prisma client extension is as simple as doing an npm install as you would in any node app.

After the imports, I re-export some library constructs so our users don’t have to care or know about any upstream libraries. Those combined with our main extension method signature mean the user can easily see all available configuration options by inspecting the type signature in their IDE.

export { Audience, OperationType, StreamdalConfigs };

export type StreamdalArgs = {
  streamdalAudience?: Audience;
};

export const streamdal = (streamdalConfigs: StreamdalConfigs) => {
...
}

With that our users can now simply wrap their existing Prisma client with this extension to get access to the added functionality and without breaking any existing use cases. The extension wrapper looks like this.

const prisma = new PrismaClient().$extends(streamdal(streamdalConfigs));

Within the extension, we immediately construct a few persistent things we need and name our extension using the convention Prisma recommends:

export const streamdal = (streamdalConfigs: StreamdalConfigs) => {
  const streamdal = new Streamdal(streamdalConfigs);
  const decoder = new TextDecoder();
  const encoder = new TextEncoder();

  return Prisma.defineExtension({
    name: "prisma-extension-streamdal",
    ...
  });
};

Next we extend the create operation on all models. Here we add the StreamdalArgs type we declared above to the existing Prisma create type so users can again inspect the create type signature to see the newly available options.

return Prisma.defineExtension({
    name: "prisma-extension-streamdal",
    model: {
      $allModels: {
        async create<T, A>(
          this: T,
          args?: Prisma.Exact<A, Prisma.Args<T, "create">> & StreamdalArgs,
        ): Promise<Prisma.Result<T, A, "create">> {
          ...
        },
      },
    },
  });

Note that the new args are optional. This allows us to add new functionality without breaking existing use cases. If the new args are present, we will wrap the operation with the new functionality, if not we simply invoke the existing operation. We do this by getting access to the parent operation via the extension context like so:

const ctx = Prisma.getExtensionContext(this);
...
return (ctx as any).$parent[ctx.$name as any].create(rest);

The exact details of the extended operation logic might vary depending on the operation. For example in our case, for read operations we execute our pipelines on the data after it fetched and for mutations we execute pipelines on the data before it is sent to the database. The idea here is that our users might want to do something like detect PII and mask it before it enters the database or leaves the system.

You can see that intervention in the example create we implemented.

const { streamdalAudience, ...rest }: any & StreamdalArgs = args || {};
const ctx = Prisma.getExtensionContext(this);

if (streamdalAudience) {
  const streamdalResult: SDKResponse = await streamdal.process({
    audience: streamdalAudience,
    data: encoder.encode(JSON.stringify(rest.data)),
  });
  const data = decoder.decode(streamdalResult.data);
  return (ctx as any).$parent[ctx.$name as any].create({
    data: JSON.parse(data),
  });
}

return (ctx as any).$parent[ctx.$name as any].create(rest);

Next, we add extended model methods for all the operations we care about with the appropriate pipeline logic before or after the underlying operation. That means method wrappers for all the creates, finds, and updates Prisma provides. See them all here.

Now that we’ve done that, we publish to NPM so all our users need to do to get started is: npm install @streamdal/prisma-extension-streamdal.

If you are curious about our Code-Native pipeline platform you can read more about it at streamdal.com or get started here.

We’ve got a community!

Want to nerd out with me and other misfits about your experiences with monorepos, deep-tech, or anything engineering-related?

Join our Discord, we’d love to have you!