@xata.io/drizzle 0.0.0-alpha.vab97e27be6032a4c2bb926be93ba2d31c905cb75 → 0.0.0-alpha.vac31b16bd60e3a0e558ad80de54b7dcad6cd8221

package/test/drizzle.test.ts

@@ -1,6 +1,6 @@
 import { BaseClient, HostProvider, parseProviderString, XataApiClient } from '@xata.io/client';
 import 'dotenv/config';
-import { desc, DrizzleError, eq, gt, gte, or, placeholder, sql, TransactionRollbackError } from 'drizzle-orm';
+import { desc, DrizzleError, eq, gt, gte, or, placeholder, sql } from 'drizzle-orm';
 import { Client } from 'pg';
 import { afterAll, afterEach, beforeAll, beforeEach, describe, expectTypeOf, test } from 'vitest';
 import { drizzle as drizzlePg, type XataDatabase } from '../src/pg';
@@ -9,7 +9,7 @@ import * as schema from './schema';
 
 const { usersTable, postsTable, commentsTable, usersToGroupsTable, groupsTable } = schema;
 
-const ENABLE_LOGGING = true;
+const ENABLE_LOGGING = false;
 
 declare module 'vitest' {
   export interface TestContext {
@@ -51,25 +51,48 @@ function getDomain(host: HostProvider) {
   }
 }
 
-beforeAll(async () => {
-  await api.database.createDatabase({
-    workspace,
-    database,
-    data: { region, branchName: 'main' },
-    headers: { 'X-Features': 'feat-pgroll-migrations=1' }
-  });
+function getDrizzleClient(type: string, branch: string) {
+  if (type === 'http') {
+    const xata = new BaseClient({
+      apiKey,
+      host,
+      clientName: 'sdk-tests',
+      databaseURL: `https://${workspace}.${region}.${getDomain(host)}/db/${database}`,
+      branch
+    });
 
-  await waitForReplication();
+    return { db: drizzleHttp(xata, { schema, logger: ENABLE_LOGGING }) };
+  } else if (type === 'pg') {
+    const client = new Client({
+      connectionString: `postgresql://${workspace}:${apiKey}@${region}.sql.${getDomain(
+        host
+      )}:5432/${database}:${branch}`,
+      ssl: true
+    });
 
-  const client = new Client({
-    connectionString: `postgresql://${workspace}:${apiKey}@${region}.sql.${getDomain(host)}:5432/${database}:main`
-  });
+    return { db: drizzlePg(client, { schema, logger: ENABLE_LOGGING }), client };
+  } else {
+    throw new Error(`Unknown type: ${type}`);
+  }
+}
+
+describe.concurrent.each([{ type: 'pg' }, { type: 'http' }])('Drizzle $type', ({ type }) => {
+  beforeAll(async () => {
+    await api.database.createDatabase({
+      workspace,
+      database,
+      data: { region, branchName: 'main' },
+      headers: { 'X-Features': 'feat-pgroll-migrations=1' }
+    });
 
-  await client.connect();
-  const db = drizzlePg(client, { schema, logger: ENABLE_LOGGING });
+    await waitForReplication();
 
-  await db.execute(
-    sql`
+    // For now, run the migrations via wire protocol
+    const { client, db } = getDrizzleClient('pg', 'main');
+    await client?.connect();
+
+    await db.execute(
+      sql`
 			CREATE TABLE "users" (
 				"id" serial PRIMARY KEY NOT NULL,
 				"name" text NOT NULL,
@@ -77,27 +100,27 @@ beforeAll(async () => {
 				"invited_by" int REFERENCES "users"("id")
 			);
 		`
-  );
-  await db.execute(
-    sql`
+    );
+    await db.execute(
+      sql`
 			CREATE TABLE IF NOT EXISTS "groups" (
 				"id" serial PRIMARY KEY NOT NULL,
 				"name" text NOT NULL,
 				"description" text
 			);
 		`
-  );
-  await db.execute(
-    sql`
+    );
+    await db.execute(
+      sql`
 			CREATE TABLE IF NOT EXISTS "users_to_groups" (
 				"id" serial PRIMARY KEY NOT NULL,
 				"user_id" int REFERENCES "users"("id"),
 				"group_id" int REFERENCES "groups"("id")
 			);
 		`
-  );
-  await db.execute(
-    sql`
+    );
+    await db.execute(
+      sql`
 			CREATE TABLE IF NOT EXISTS "posts" (
 				"id" serial PRIMARY KEY NOT NULL,
 				"content" text NOT NULL,
@@ -105,9 +128,9 @@ beforeAll(async () => {
 				"created_at" timestamp with time zone DEFAULT now() NOT NULL
 			);
 		`
-  );
-  await db.execute(
-    sql`
+    );
+    await db.execute(
+      sql`
 			CREATE TABLE IF NOT EXISTS "comments" (
 				"id" serial PRIMARY KEY NOT NULL,
 				"content" text NOT NULL,
@@ -116,9 +139,9 @@ beforeAll(async () => {
 				"created_at" timestamp with time zone DEFAULT now() NOT NULL
 			);
 		`
-  );
-  await db.execute(
-    sql`
+    );
+    await db.execute(
+      sql`
 			CREATE TABLE IF NOT EXISTS "comment_likes" (
 				"id" serial PRIMARY KEY NOT NULL,
 				"creator" int REFERENCES "users"("id"),
@@ -126,42 +149,24 @@ beforeAll(async () => {
 				"created_at" timestamp with time zone DEFAULT now() NOT NULL
 			);
 		`
-  );
+    );
 
-  await client.end();
-});
+    await client?.end();
+  });
 
-afterAll(async () => {
-  await api.database.deleteDatabase({ workspace, database });
-});
+  afterAll(async () => {
+    await api.database.deleteDatabase({ workspace, database });
+  });
 
-describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type', ({ type }) => {
   beforeEach(async (ctx) => {
     ctx.branch = `test-${Math.random().toString(36).substring(7)}`;
     await api.branches.createBranch({ workspace, database, region, branch: ctx.branch, from: 'main' });
 
-    if (type === 'http') {
-      const xata = new BaseClient({
-        apiKey,
-        host,
-        clientName: 'sdk-tests',
-        databaseURL: `https://${workspace}.${region}.${getDomain(host)}/db/${database}`,
-        branch: ctx.branch
-      });
-
-      ctx.db = drizzleHttp(xata, { schema, logger: ENABLE_LOGGING });
-    } else if (type === 'pg') {
-      ctx.client = new Client({
-        connectionString: `postgresql://${workspace}:${apiKey}@${region}.sql.${getDomain(host)}:5432/${database}:${
-          ctx.branch
-        }`
-      });
-
-      await ctx.client.connect();
-      ctx.db = drizzlePg(ctx.client, { schema, logger: ENABLE_LOGGING });
-    } else {
-      throw new Error(`Unknown type: ${type}`);
-    }
+    const { db, client } = getDrizzleClient(type, ctx.branch);
+    await client?.connect();
+
+    ctx.db = db;
+    ctx.client = client;
   });
 
   afterEach(async (ctx) => {
@@ -886,133 +891,6 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
     });
   });
 
-  test('[Find Many] Get users with posts in transaction', async (ctx) => {
-    let usersWithPosts: {
-      id: number;
-      name: string;
-      verified: boolean;
-      invitedBy: number | null;
-      posts: {
-        id: number;
-        content: string;
-        ownerId: number | null;
-        createdAt: Date;
-      }[];
-    }[] = [];
-
-    await ctx.db.transaction(async (tx) => {
-      await tx.insert(usersTable).values([
-        { id: 1, name: 'Dan' },
-        { id: 2, name: 'Andrew' },
-        { id: 3, name: 'Alex' }
-      ]);
-
-      await tx.insert(postsTable).values([
-        { ownerId: 1, content: 'Post1' },
-        { ownerId: 1, content: 'Post1.1' },
-        { ownerId: 2, content: 'Post2' },
-        { ownerId: 3, content: 'Post3' }
-      ]);
-
-      usersWithPosts = await tx.query.usersTable.findMany({
-        where: ({ id }, { eq }) => eq(id, 1),
-        with: {
-          posts: {
-            where: ({ id }, { eq }) => eq(id, 1)
-          }
-        }
-      });
-    });
-
-    expectTypeOf(usersWithPosts).toEqualTypeOf<
-      {
-        id: number;
-        name: string;
-        verified: boolean;
-        invitedBy: number | null;
-        posts: {
-          id: number;
-          content: string;
-          ownerId: number | null;
-          createdAt: Date;
-        }[];
-      }[]
-    >();
-
-    ctx.expect(usersWithPosts.length).eq(1);
-    ctx.expect(usersWithPosts[0]?.posts.length).eq(1);
-
-    ctx.expect(usersWithPosts[0]).toEqual({
-      id: 1,
-      name: 'Dan',
-      verified: false,
-      invitedBy: null,
-      posts: [{ id: 1, ownerId: 1, content: 'Post1', createdAt: usersWithPosts[0]?.posts[0]?.createdAt }]
-    });
-  });
-
-  test('[Find Many] Get users with posts in rollbacked transaction', async (ctx) => {
-    let usersWithPosts: {
-      id: number;
-      name: string;
-      verified: boolean;
-      invitedBy: number | null;
-      posts: {
-        id: number;
-        content: string;
-        ownerId: number | null;
-        createdAt: Date;
-      }[];
-    }[] = [];
-
-    await ctx
-      .expect(
-        ctx.db.transaction(async (tx) => {
-          await tx.insert(usersTable).values([
-            { id: 1, name: 'Dan' },
-            { id: 2, name: 'Andrew' },
-            { id: 3, name: 'Alex' }
-          ]);
-
-          await tx.insert(postsTable).values([
-            { ownerId: 1, content: 'Post1' },
-            { ownerId: 1, content: 'Post1.1' },
-            { ownerId: 2, content: 'Post2' },
-            { ownerId: 3, content: 'Post3' }
-          ]);
-
-          tx.rollback();
-
-          usersWithPosts = await tx.query.usersTable.findMany({
-            where: ({ id }, { eq }) => eq(id, 1),
-            with: {
-              posts: {
-                where: ({ id }, { eq }) => eq(id, 1)
-              }
-            }
-          });
-        })
-      )
-      .rejects.toThrowError(new TransactionRollbackError());
-
-    expectTypeOf(usersWithPosts).toEqualTypeOf<
-      {
-        id: number;
-        name: string;
-        verified: boolean;
-        invitedBy: number | null;
-        posts: {
-          id: number;
-          content: string;
-          ownerId: number | null;
-          createdAt: Date;
-        }[];
-      }[]
-    >();
-
-    ctx.expect(usersWithPosts.length).eq(0);
-  });
-
   // select only custom
   test('[Find Many] Get only custom fields', async (ctx) => {
     await ctx.db.insert(usersTable).values([