[Drizzle] Add support for XataFiles

.changeset/yellow-mugs-fry.md

@@ -0,0 +1,5 @@
+---
+'@xata.io/drizzle': patch
+---
+
+Update drizzle driver

packages/plugin-client-drizzle/package.json

@@ -10,6 +10,11 @@
       "import": "./dist/index.mjs",
       "require": "./dist/index.cjs",
       "types": "./dist/index.d.ts"
+    },
+    "./pg": {
+      "import": "./dist/pg.mjs",
+      "require": "./dist/pg.cjs",
+      "types": "./dist/pg.d.ts"
     }
   },
   "scripts": {
@@ -22,10 +27,12 @@
     "url": "https://github.com/xataio/client-ts/issues"
   },
   "dependencies": {
-    "@xata.io/client": "workspace:*"
+    "@xata.io/client": "workspace:*",
+    "pg": "^8.11.3"
   },
   "devDependencies": {
-    "drizzle-orm": "^0.29.3"
+    "@types/pg": "^8.11.2",
+    "drizzle-orm": "^0.29.5"
   },
   "peerDependencies": {
     "drizzle-orm": "^0.28.5"

packages/plugin-client-drizzle/rollup.config.mjs

@@ -25,5 +25,29 @@ export default [
       file: `dist/index.d.ts`,
       format: 'es'
     }
-  }
+  },
+  {
+    input: 'src/pg/index.ts',
+    plugins: [esbuild()],
+    output: [
+      {
+        file: `dist/pg.cjs`,
+        format: 'cjs',
+        sourcemap: true
+      },
+      {
+        file: `dist/pg.mjs`,
+        format: 'es',
+        sourcemap: true
+      }
+    ]
+  },
+  {
+    input: 'src/pg/index.ts',
+    plugins: [dts()],
+    output: {
+      file: `dist/pg.d.ts`,
+      format: 'es'
+    }
+  },
 ];

packages/plugin-client-drizzle/src/http/driver.ts

@@ -0,0 +1,79 @@
+import {
+  DefaultLogger,
+  DrizzleConfig,
+  ExtractTablesWithRelations,
+  Logger,
+  RelationalSchemaConfig,
+  TablesRelationalConfig,
+  createTableRelationsHelpers,
+  entityKind,
+  extractTablesRelationalConfig
+} from 'drizzle-orm';
+import { XataHttpSession, type XataHttpClient, type XataHttpQueryResultHKT } from './session.js';
+import { PgDatabase, PgDialect } from 'drizzle-orm/pg-core';
+
+export interface XataDriverOptions {
+  logger?: Logger;
+}
+
+export class XataHttpDriver {
+  static readonly [entityKind]: string = 'XataDriver';
+
+  constructor(private client: XataHttpClient, private dialect: PgDialect, private options: XataDriverOptions = {}) {
+    this.initMappers();
+  }
+
+  createSession(
+    schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined
+  ): XataHttpSession<Record<string, unknown>, TablesRelationalConfig> {
+    return new XataHttpSession(this.client, this.dialect, schema, {
+      logger: this.options.logger
+    });
+  }
+
+  initMappers() {
+    // TODO: Add custom type parsers
+  }
+}
+
+export class XataHttpDatabase<TSchema extends Record<string, unknown> = Record<string, never>> extends PgDatabase<
+  XataHttpQueryResultHKT,
+  TSchema
+> {
+  static readonly [entityKind]: string = 'XataHttpDatabase';
+
+  /** @internal */
+  declare readonly session: XataHttpSession<TSchema, ExtractTablesWithRelations<TSchema>>;
+}
+
+export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
+  client: XataHttpClient,
+  config: DrizzleConfig<TSchema> = {}
+): XataHttpDatabase<TSchema> {
+  const dialect = new PgDialect();
+  let logger;
+  if (config.logger === true) {
+    logger = new DefaultLogger();
+  } else if (config.logger !== false) {
+    logger = config.logger;
+  }
+
+  let schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined;
+  if (config.schema) {
+    const tablesConfig = extractTablesRelationalConfig(config.schema, createTableRelationsHelpers);
+    schema = {
+      fullSchema: config.schema,
+      schema: tablesConfig.tables,
+      tableNamesMap: tablesConfig.tableNamesMap
+    };
+  }
+
+  const driver = new XataHttpDriver(client, dialect, { logger });
+  const session = driver.createSession(schema);
+
+  return new XataHttpDatabase(
+    dialect,
+    session,
+    schema as RelationalSchemaConfig<ExtractTablesWithRelations<TSchema>> | undefined
+  );
+}

packages/plugin-client-drizzle/src/http/index.ts

@@ -0,0 +1,2 @@
+export * from './driver.js';
+export * from './session.js';

packages/plugin-client-drizzle/src/http/migrator.ts

@@ -0,0 +1,60 @@
+import { MigrationConfig, readMigrationFiles } from 'drizzle-orm/migrator';
+import { XataHttpDatabase } from './driver';
+import { SQL, sql } from 'drizzle-orm';
+
+/**
+ * This function reads migrationFolder and execute each unapplied migration and mark it as executed in database
+ *
+ * NOTE: The Xata HTTP driver does not support transactions. This means that if any part of a migration fails,
+ * no rollback will be executed. Currently, you will need to handle unsuccessful migration yourself.
+ * @param db - drizzle db instance
+ * @param config - path to migration folder generated by drizzle-kit
+ */
+export async function migrate<TSchema extends Record<string, unknown>>(
+  db: XataHttpDatabase<TSchema>,
+  config: string | MigrationConfig
+) {
+  const migrations = readMigrationFiles(config);
+  const migrationsTable =
+    typeof config === 'string' ? '__drizzle_migrations' : config.migrationsTable ?? '__drizzle_migrations';
+  const migrationsSchema = typeof config === 'string' ? 'drizzle' : config.migrationsSchema ?? 'drizzle';
+  const migrationTableCreate = sql`
+		CREATE TABLE IF NOT EXISTS ${sql.identifier(migrationsSchema)}.${sql.identifier(migrationsTable)} (
+			id SERIAL PRIMARY KEY,
+			hash text NOT NULL,
+			created_at bigint
+		)
+	`;
+  await db.session.execute(sql`CREATE SCHEMA IF NOT EXISTS ${sql.identifier(migrationsSchema)}`);
+  await db.session.execute(migrationTableCreate);
+
+  const dbMigrations = await db.session.all<{
+    id: number;
+    hash: string;
+    created_at: string;
+  }>(
+    sql`select id, hash, created_at from ${sql.identifier(migrationsSchema)}.${sql.identifier(
+      migrationsTable
+    )} order by created_at desc limit 1`
+  );
+
+  const lastDbMigration = dbMigrations[0];
+  const rowsToInsert: SQL[] = [];
+  for await (const migration of migrations) {
+    if (!lastDbMigration || Number(lastDbMigration.created_at) < migration.folderMillis) {
+      for (const stmt of migration.sql) {
+        await db.session.execute(sql.raw(stmt));
+      }
+
+      rowsToInsert.push(
+        sql`insert into ${sql.identifier(migrationsSchema)}.${sql.identifier(
+          migrationsTable
+        )} ("hash", "created_at") values(${migration.hash}, ${migration.folderMillis})`
+      );
+    }
+  }
+
+  for await (const rowToInsert of rowsToInsert) {
+    await db.session.execute(rowToInsert);
+  }
+}