drizzle-driver: add smaller Drizzle beta support

Author: ziahamza

packages/drizzle-driver/README.md

@@ -6,6 +6,28 @@ This package (`@powersync/drizzle-driver`) brings the benefits of an ORM through
 
 The `drizzle-driver` package is currently in an Beta release.
 
+## Drizzle Version Support
+
+- `drizzle-orm@0.44.x`
+  - Pass `schema`.
+  - `db.query` keeps working as before.
+  - `db._query` is available as an alias to the same legacy relational query builder.
+- `drizzle-orm@1.0.0-beta`
+  - Pass `relations`.
+  - Use `db.query` for relational queries.
+  - If you only pass `schema`, the driver exposes a basic `db.query.<table>.findMany()` surface for simple table queries, but relational `with` queries should use `relations`.
+
+## Drizzle Version Support
+
+- `drizzle-orm@0.44.x`
+  - Use `schema`.
+  - `db.query` continues to use the legacy relational query builder.
+  - `db._query` is available as a compatibility alias to the same legacy query builder.
+- `drizzle-orm@1.0.0-beta`
+  - Use `relations` for the new relational query builder and query via `db.query`.
+  - Keep using `schema` if you want the legacy relational query builder on `db._query`.
+  - If you only pass `schema`, the driver synthesizes a basic `db.query.<table>.findMany()` surface for simple table queries.
+
 ## Getting Started
 
 Set up the PowerSync Database and wrap it with Drizzle.
@@ -67,6 +89,112 @@ export const db = wrapPowerSyncWithDrizzle(powerSyncDb, {
 });
 ```
 
+### Drizzle v1 beta
+
+For Drizzle beta relational queries, pass `relations` and use `db.query`.
+
+```js
+import { defineRelations } from 'drizzle-orm';
+import { wrapPowerSyncWithDrizzle } from '@powersync/drizzle-driver';
+
+export const relations = defineRelations({ lists, todos }, (r) => ({
+  lists: {
+    todos: r.many.todos({
+      from: r.lists.id,
+      to: r.todos.list_id
+    })
+  },
+  todos: {
+    list: r.one.lists({
+      from: r.todos.list_id,
+      to: r.lists.id,
+      optional: false
+    })
+  }
+}));
+
+export const db = wrapPowerSyncWithDrizzle(powerSyncDb, {
+  relations
+});
+
+const listsWithTodos = await db.query.lists.findMany({
+  with: {
+    todos: true
+  }
+});
+```
+
+### Drizzle v1 beta
+
+For Drizzle beta relational queries, pass `relations` and use `db.query`.
+
+```js
+import { defineRelations } from 'drizzle-orm';
+import { wrapPowerSyncWithDrizzle } from '@powersync/drizzle-driver';
+
+export const relations = defineRelations({ lists, todos }, (r) => ({
+  lists: {
+    todos: r.many.todos({
+      from: r.lists.id,
+      to: r.todos.list_id
+    })
+  },
+  todos: {
+    list: r.one.lists({
+      from: r.todos.list_id,
+      to: r.lists.id,
+      optional: false
+    })
+  }
+}));
+
+export const db = wrapPowerSyncWithDrizzle(powerSyncDb, {
+  relations
+});
+
+const listsWithTodos = await db.query.lists.findMany({
+  with: {
+    todos: true
+  }
+});
+```
+
+### Partial upgrade on Drizzle beta
+
+If you are upgrading incrementally, keep legacy relation definitions in `drizzle-orm/_relations`, pass them via `schema`, and use `db._query`.
+
+```js
+import { relations } from 'drizzle-orm/_relations';
+
+const listsRelations = relations(lists, ({ many }) => ({
+  todos: many(todos)
+}));
+
+const todosRelations = relations(todos, ({ one }) => ({
+  list: one(lists, {
+    fields: [todos.list_id],
+    references: [lists.id]
+  })
+}));
+
+const legacySchema = {
+  lists,
+  todos,
+  listsRelations,
+  todosRelations
+};
+
+export const db = wrapPowerSyncWithDrizzle(powerSyncDb, {
+  schema: legacySchema
+});
+
+const listsWithTodos = await db._query.lists.findMany({
+  with: {
+    todos: true
+  }
+});
+```
+
 ## Schema Conversion
 
 The `DrizzleAppSchema` constructor simplifies the process of integrating Drizzle with PowerSync. It infers the local [PowerSync schema](https://docs.powersync.com/installation/client-side-setup/define-your-schema) from your Drizzle schema definition, providing a unified development experience.

packages/drizzle-driver/package.json

@@ -47,7 +47,7 @@
   },
   "peerDependencies": {
     "@powersync/common": "workspace:^1.50.0",
-    "drizzle-orm": "<1.0.0"
+    "drizzle-orm": ">=0.44.7 <1.0.0 || >=1.0.0-beta.1 <1.0.0"
   },
   "devDependencies": {
     "@journeyapps/wa-sqlite": "catalog:",

packages/drizzle-driver/src/index.ts

@@ -1,6 +1,7 @@
 import {
   wrapPowerSyncWithDrizzle,
   type DrizzleQuery,
+  type PowerSyncDrizzleConfig,
   type PowerSyncSQLiteDatabase
 } from './sqlite/PowerSyncSQLiteDatabase.js';
 import { toCompilableQuery } from './utils/compilableQuery.js';
@@ -24,6 +25,7 @@ export {
   DrizzleTableWithPowerSyncOptions,
   Expand,
   ExtractPowerSyncColumns,
+  PowerSyncDrizzleConfig,
   PowerSyncSQLiteDatabase,
   TableName,
   TablesFromSchemaEntries,

packages/drizzle-driver/src/sqlite/PowerSyncSQLiteBaseSession.ts

@@ -16,6 +16,9 @@ import {
 } from 'drizzle-orm/sqlite-core/session';
 import { PowerSyncSQLitePreparedQuery, type ContextProvider } from './PowerSyncSQLitePreparedQuery.js';
 
+const SQLiteSessionBase = SQLiteSession as unknown as abstract new (...args: any[]) => any;
+const SQLiteTransactionBase = SQLiteTransaction as unknown as abstract new (...args: any[]) => any;
+
 export interface PowerSyncSQLiteSessionOptions {
   logger?: Logger;
 }
@@ -27,14 +30,14 @@ export type PowerSyncSQLiteTransactionConfig = SQLiteTransactionConfig & {
 export class PowerSyncSQLiteTransaction<
   TFullSchema extends Record<string, unknown>,
   TSchema extends TablesRelationalConfig
-> extends SQLiteTransaction<'async', QueryResult, TFullSchema, TSchema> {
+> extends SQLiteTransactionBase {
   static readonly [entityKind]: string = 'PowerSyncSQLiteTransaction';
 }
 
 export class PowerSyncSQLiteBaseSession<
   TFullSchema extends Record<string, unknown>,
   TSchema extends TablesRelationalConfig
-> extends SQLiteSession<'async', QueryResult, TFullSchema, TSchema> {
+> extends SQLiteSessionBase {
   static readonly [entityKind]: string = 'PowerSyncSQLiteBaseSession';
 
   protected logger: Logger;
@@ -43,7 +46,8 @@ export class PowerSyncSQLiteBaseSession<
     protected contextProvider: ContextProvider,
     protected dialect: SQLiteAsyncDialect,
     protected schema: RelationalSchemaConfig<TSchema> | undefined,
-    protected options: PowerSyncSQLiteSessionOptions = {}
+    protected options: PowerSyncSQLiteSessionOptions = {},
+    protected relations: Record<string, unknown> = {}
   ) {
     super(dialect);
     this.logger = options.logger ?? new NoopLogger();
@@ -75,6 +79,27 @@ export class PowerSyncSQLiteBaseSession<
     );
   }
 
+  prepareRelationalQuery<T extends PreparedQueryConfigBase & { type: 'async' }>(
+    query: Query,
+    fields: SelectedFieldsOrdered | undefined,
+    executeMethod: SQLiteExecuteMethod,
+    customResultMapper: (rows: Record<string, unknown>[], mapColumnValue?: (value: unknown) => unknown) => unknown
+  ): PowerSyncSQLitePreparedQuery<T> {
+    return new PowerSyncSQLitePreparedQuery(
+      this.contextProvider,
+      query,
+      this.logger,
+      fields,
+      executeMethod,
+      false,
+      customResultMapper,
+      undefined,
+      { type: 'select', tables: [] },
+      undefined,
+      true
+    );
+  }
+
   transaction<T>(
     _transaction: (tx: PowerSyncSQLiteTransaction<TFullSchema, TSchema>) => T,
     _config: PowerSyncSQLiteTransactionConfig = {}

packages/drizzle-driver/src/sqlite/PowerSyncSQLiteDatabase.ts

@@ -8,31 +8,53 @@ import {
 import { Query } from 'drizzle-orm';
 import { DefaultLogger } from 'drizzle-orm/logger';
 import {
-  createTableRelationsHelpers,
-  extractTablesRelationalConfig,
   ExtractTablesWithRelations,
   TableRelationalConfig,
   type RelationalSchemaConfig,
   type TablesRelationalConfig
 } from 'drizzle-orm/relations';
+import * as DrizzleRelations from 'drizzle-orm/relations';
 import { SQLiteSession, SQLiteTable, SQLiteTransaction } from 'drizzle-orm/sqlite-core';
 import { BaseSQLiteDatabase } from 'drizzle-orm/sqlite-core/db';
 import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core/dialect';
 import { RelationalQueryBuilder } from 'drizzle-orm/sqlite-core/query-builders/query';
 import type { DrizzleConfig } from 'drizzle-orm/utils';
+import { extractFallbackRelations, isDrizzleBetaRuntime } from '../utils/drizzleCompat.js';
 import { toCompilableQuery } from './../utils/compilableQuery.js';
 import { PowerSyncSQLiteBaseSession, PowerSyncSQLiteTransactionConfig } from './PowerSyncSQLiteBaseSession.js';
 import { PowerSyncSQLiteSession } from './PowerSyncSQLiteSession.js';
 
+const BaseSQLiteDatabaseBase = BaseSQLiteDatabase as unknown as abstract new (...args: any[]) => any;
+
 export type DrizzleQuery<T> = { toSQL(): Query; execute(): Promise<T | T[]> };
 
+export type PowerSyncDrizzleConfig<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+  TRelations extends Record<string, unknown> = Record<string, never>
+> = DrizzleConfig<TSchema> & {
+  relations?: TRelations;
+};
+
 export class PowerSyncSQLiteDatabase<
-  TSchema extends Record<string, unknown> = Record<string, never>
-> extends BaseSQLiteDatabase<'async', QueryResult, TSchema> {
+  TSchema extends Record<string, unknown> = Record<string, never>,
+  TRelations extends Record<string, unknown> = Record<string, never>
+> extends BaseSQLiteDatabaseBase {
+  declare _:
+    | {
+        schema: TablesRelationalConfig | undefined;
+        fullSchema: TSchema;
+        tableNamesMap: Record<string, string>;
+      }
+    | undefined;
+
+  declare query: Record<string, unknown>;
+  declare _query?: Record<string, unknown>;
+
   private db: AbstractPowerSyncDatabase;
 
-  constructor(db: AbstractPowerSyncDatabase, config: DrizzleConfig<TSchema> = {}) {
+  constructor(db: AbstractPowerSyncDatabase, config: PowerSyncDrizzleConfig<TSchema, TRelations> = {}) {
     const dialect = new SQLiteAsyncDialect({ casing: config.casing });
+    const drizzleBetaRuntime = isDrizzleBetaRuntime(dialect);
     let logger;
     if (config.logger === true) {
       logger = new DefaultLogger();
@@ -41,40 +63,47 @@ export class PowerSyncSQLiteDatabase<
     }
 
     let schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined;
-    if (config.schema) {
-      const tablesConfig = extractTablesRelationalConfig(config.schema, createTableRelationsHelpers);
+    if (config.schema && !drizzleBetaRuntime) {
+      const tablesConfig = (DrizzleRelations as any).extractTablesRelationalConfig(
+        config.schema,
+        (DrizzleRelations as any).createTableRelationsHelpers
+      );
       schema = {
         fullSchema: config.schema,
         schema: tablesConfig.tables,
         tableNamesMap: tablesConfig.tableNamesMap
       };
     }
 
-    const session = new PowerSyncSQLiteSession(db, dialect, schema, {
-      logger
-    });
+    const relations = (config.relations ??
+      (drizzleBetaRuntime ? extractFallbackRelations(config.schema) : {})) as Record<string, unknown>;
 
-    super('async', dialect, session as any, schema as any);
+    const session = new PowerSyncSQLiteSession(db, dialect, schema, { logger }, relations);
+
+    if (drizzleBetaRuntime) {
+      super('async', dialect, session as any, relations as any, undefined, undefined, true);
+    } else {
+      super('async', dialect, session as any, schema as any);
+    }
     this.db = db;
 
-    /**
-     * A hack in order to use read locks for `db.query.users.findMany()` etc queries.
-     * We don't currently get queryMetadata for these queries, so we can't use the regular session.
-     * This session always uses read locks.
-     */
-    const querySession = new PowerSyncSQLiteBaseSession(
-      {
-        useReadContext: (callback) => db.readLock(callback),
-        useWriteContext: (callback) => db.readLock(callback)
-      },
-      dialect,
-      schema,
-      {
-        logger
-      }
-    );
-    if (this._.schema) {
-      // https://github.com/drizzle-team/drizzle-orm/blob/ad4ddd444d066b339ffd5765cb6ec3bf49380189/drizzle-orm/src/sqlite-core/db.ts#L72
+    if (!drizzleBetaRuntime && this._?.schema) {
+      /**
+       * A hack in order to use read locks for `db.query.users.findMany()` etc queries.
+       * We don't currently get queryMetadata for these queries, so we can't use the regular session.
+       * This session always uses read locks.
+       */
+      const querySession = new PowerSyncSQLiteBaseSession(
+        {
+          useReadContext: (callback) => db.readLock(callback),
+          useWriteContext: (callback) => db.readLock(callback)
+        },
+        dialect,
+        schema,
+        {
+          logger
+        }
+      );
       const query = this.query as {
         [K in keyof TSchema]: RelationalQueryBuilder<'async', any, any, any>;
       };
@@ -87,9 +116,11 @@ export class PowerSyncSQLiteDatabase<
           schema!.fullSchema[tableName] as SQLiteTable,
           columns as TableRelationalConfig,
           dialect,
-          querySession as SQLiteSession<'async', any, any, any>
+          querySession as unknown as SQLiteSession<'async', any, any, any>
         );
       }
+
+      this._query = this.query;
     }
   }
 
@@ -107,9 +138,12 @@ export class PowerSyncSQLiteDatabase<
   }
 }
 
-export function wrapPowerSyncWithDrizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
+export function wrapPowerSyncWithDrizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+  TRelations extends Record<string, unknown> = Record<string, never>
+>(
   db: AbstractPowerSyncDatabase,
-  config: DrizzleConfig<TSchema> = {}
-): PowerSyncSQLiteDatabase<TSchema> {
-  return new PowerSyncSQLiteDatabase<TSchema>(db, config);
+  config: PowerSyncDrizzleConfig<TSchema, TRelations> = {}
+): PowerSyncSQLiteDatabase<TSchema, TRelations> {
+  return new PowerSyncSQLiteDatabase<TSchema, TRelations>(db, config);
 }