@xata.io/drizzle 0.0.0-alpha.vcf8f11cade1aeb5e318fc16f42ddc96e46b17b8e → 0.0.0-alpha.vcfd6a67de8fbe1990b37c436166fb2bdce576c79

package/test/{pg.test.ts → drizzle.test.ts}

@@ -1,6 +1,6 @@
 import { BaseClient, HostProvider, parseProviderString, XataApiClient } from '@xata.io/client';
 import 'dotenv/config';
-import { desc, 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 {
@@ -25,10 +25,13 @@ if (apiKey === '') throw new Error('XATA_API_KEY environment variable is not set
 const workspace = (process.env.XATA_WORKSPACE ?? '').split('-').pop() ?? '';
 if (workspace === '') throw new Error('XATA_WORKSPACE environment variable is not set');
 
-const region = process.env.XATA_REGION || 'eu-west-1';
-
 const host = parseProviderString(process.env.XATA_API_PROVIDER) ?? 'production';
 
+// TODO: Branches for pgroll only work in some regions for now
+// const region = process.env.XATA_REGION || 'us-east-1';
+const region =
+  host === 'production' ? 'us-east-1' : host === 'staging' ? 'eu-west-1' : process.env.XATA_REGION || 'us-east-1';
+
 const database = `drizzle-test-${Math.random().toString(36).substring(7)}`;
 
 const api = new XataApiClient({ apiKey, host, clientName: 'sdk-tests' });
@@ -48,27 +51,49 @@ 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, database: 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`,
-    // Not sure why, but we are getting `error: SSL required` sometimes
-    ssl: { rejectUnauthorized: false }
-  });
+    return { db: drizzlePg(client, { schema, logger: ENABLE_LOGGING }), client };
+  } else {
+    throw new Error(`Unknown type: ${type}`);
+  }
+}
 
-  await client.connect();
-  const db = drizzlePg(client, { schema, logger: ENABLE_LOGGING });
+describe.concurrent.each([{ type: 'pg' }, { type: 'http' }])('Drizzle $type', ({ type }) => {
+  const dbName = `${database}-${type}`;
 
-  await db.execute(
-    sql`
+  beforeAll(async () => {
+    await api.databases.createDatabase({
+      pathParams: { workspaceId: workspace, dbName },
+      body: { region, branchName: 'main' },
+      headers: { 'X-Features': 'feat-pgroll-migrations=1' }
+    });
+
+    await waitForReplication(dbName);
+
+    // For now, run the migrations via wire protocol
+    const { client, db } = getDrizzleClient('pg', dbName, 'main');
+    await client?.connect();
+
+    await db.execute(
+      sql`
 			CREATE TABLE "users" (
 				"id" serial PRIMARY KEY NOT NULL,
 				"name" text NOT NULL,
@@ -76,27 +101,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,
@@ -104,9 +129,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,
@@ -115,9 +140,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"),
@@ -125,49 +150,32 @@ 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.databases.deleteDatabase({ pathParams: { workspaceId: workspace, dbName } });
+  });
 
-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
-        }`,
-        // Not sure why, but we are getting `error: SSL required` sometimes
-        ssl: { rejectUnauthorized: false }
-      });
-
-      await ctx.client.connect();
-      ctx.db = drizzlePg(ctx.client, { schema, logger: ENABLE_LOGGING });
-    } else {
-      throw new Error(`Unknown type: ${type}`);
-    }
+    await api.branch.createBranch({
+      pathParams: { workspace, region, dbBranchName: `${dbName}:${ctx.branch}` },
+      body: { from: 'main' }
+    });
+
+    const { db, client } = getDrizzleClient(type, dbName, ctx.branch);
+    await client?.connect();
+
+    ctx.db = db;
+    ctx.client = client;
   });
 
   afterEach(async (ctx) => {
     await ctx.client?.end();
-    await api.branches.deleteBranch({ workspace, database, region, branch: ctx.branch });
+    await api.branch.deleteBranch({ pathParams: { workspace, region, dbBranchName: `${dbName}:${ctx.branch}` } });
   });
 
   /*
@@ -887,133 +895,6 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
     });
   });
 
-  test.skip('[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.skip('[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([
@@ -1330,7 +1211,7 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
           columns: {},
           where: gte(postsTable.id, 2),
           extras: ({ content }) => ({
-            lowerName: sql<string>`lower(${content})`.as('content_lower')
+            contentLower: sql<string>`lower(${content})`.as('content_lower')
           })
         }
       },
@@ -1344,7 +1225,7 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
       | {
           lowerName: string;
           posts: {
-            lowerName: string;
+            contentLower: string;
           }[];
         }
       | undefined
@@ -1354,7 +1235,7 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
 
     ctx.expect(usersWithPosts).toEqual({
       lowerName: 'dan',
-      posts: [{ lowerName: 'post1.2' }, { lowerName: 'post1.3' }]
+      posts: [{ contentLower: 'post1.2' }, { contentLower: 'post1.3' }]
     });
   });
 
@@ -1465,41 +1346,43 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
   });
 
   // columns {}
-  test.skip('[Find Many] Get select {}', async (ctx) => {
+  test('[Find Many] Get select {}', async (ctx) => {
     await ctx.db.insert(usersTable).values([
       { id: 1, name: 'Dan' },
       { id: 2, name: 'Andrew' },
       { id: 3, name: 'Alex' }
     ]);
 
-    const users = await ctx.db.query.usersTable.findMany({
-      columns: {}
-    });
-
-    ctx.expect(users.length).toBe(3);
-
-    ctx.expect(users[0]).toEqual({});
-    ctx.expect(users[1]).toEqual({});
-    ctx.expect(users[2]).toEqual({});
+    await ctx
+      .expect(
+        async () =>
+          await ctx.db.query.usersTable.findMany({
+            columns: {}
+          })
+      )
+      .rejects.toThrow(DrizzleError);
   });
 
   // columns {}
-  test.skip('[Find One] Get select {}', async (ctx) => {
+  test('[Find One] Get select {}', async (ctx) => {
     await ctx.db.insert(usersTable).values([
       { id: 1, name: 'Dan' },
       { id: 2, name: 'Andrew' },
       { id: 3, name: 'Alex' }
     ]);
 
-    const users = await ctx.db.query.usersTable.findFirst({
-      columns: {}
-    });
-
-    ctx.expect(users).toEqual({});
+    await ctx
+      .expect(
+        async () =>
+          await ctx.db.query.usersTable.findFirst({
+            columns: {}
+          })
+      )
+      .rejects.toThrow(DrizzleError);
   });
 
   // deep select {}
-  test.skip('[Find Many] Get deep select {}', async (ctx) => {
+  test('[Find Many] Get deep select {}', async (ctx) => {
     await ctx.db.insert(usersTable).values([
       { id: 1, name: 'Dan' },
       { id: 2, name: 'Andrew' },
@@ -1512,24 +1395,23 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
       { ownerId: 3, content: 'Post3' }
     ]);
 
-    const users = await ctx.db.query.usersTable.findMany({
-      columns: {},
-      with: {
-        posts: {
-          columns: {}
-        }
-      }
-    });
-
-    ctx.expect(users.length).toBe(3);
-
-    ctx.expect(users[0]).toEqual({ posts: [{}] });
-    ctx.expect(users[1]).toEqual({ posts: [{}] });
-    ctx.expect(users[2]).toEqual({ posts: [{}] });
+    await ctx
+      .expect(
+        async () =>
+          await ctx.db.query.usersTable.findMany({
+            columns: {},
+            with: {
+              posts: {
+                columns: {}
+              }
+            }
+          })
+      )
+      .rejects.toThrow(DrizzleError);
   });
 
   // deep select {}
-  test.skip('[Find One] Get deep select {}', async (ctx) => {
+  test('[Find One] Get deep select {}', async (ctx) => {
     await ctx.db.insert(usersTable).values([
       { id: 1, name: 'Dan' },
       { id: 2, name: 'Andrew' },
@@ -1542,16 +1424,19 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
       { ownerId: 3, content: 'Post3' }
     ]);
 
-    const users = await ctx.db.query.usersTable.findFirst({
-      columns: {},
-      with: {
-        posts: {
-          columns: {}
-        }
-      }
-    });
-
-    ctx.expect(users).toEqual({ posts: [{}] });
+    await ctx
+      .expect(
+        async () =>
+          await ctx.db.query.usersTable.findFirst({
+            columns: {},
+            with: {
+              posts: {
+                columns: {}
+              }
+            }
+          })
+      )
+      .rejects.toThrow(DrizzleError);
   });
 
   /*
@@ -4116,14 +4001,118 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
     });
   });
 
+  test('Get user with posts and posts with comments and comments with owner where exists', async (ctx) => {
+    await ctx.db.insert(usersTable).values([
+      { id: 1, name: 'Dan' },
+      { id: 2, name: 'Andrew' },
+      { id: 3, name: 'Alex' }
+    ]);
+
+    await ctx.db.insert(postsTable).values([
+      { id: 1, ownerId: 1, content: 'Post1' },
+      { id: 2, ownerId: 2, content: 'Post2' },
+      { id: 3, ownerId: 3, content: 'Post3' }
+    ]);
+
+    await ctx.db.insert(commentsTable).values([
+      { postId: 1, content: 'Comment1', creator: 2 },
+      { postId: 2, content: 'Comment2', creator: 2 },
+      { postId: 3, content: 'Comment3', creator: 3 }
+    ]);
+
+    const response = await ctx.db.query.usersTable.findMany({
+      with: {
+        posts: {
+          with: {
+            comments: {
+              with: {
+                author: true
+              }
+            }
+          }
+        }
+      },
+      where: (table, { exists, eq }) =>
+        exists(
+          ctx.db
+            .select({ one: sql`1` })
+            .from(usersTable)
+            .where(eq(sql`1`, table.id))
+        )
+    });
+
+    expectTypeOf(response).toEqualTypeOf<
+      {
+        id: number;
+        name: string;
+        verified: boolean;
+        invitedBy: number | null;
+        posts: {
+          id: number;
+          content: string;
+          ownerId: number | null;
+          createdAt: Date;
+          comments: {
+            id: number;
+            content: string;
+            createdAt: Date;
+            creator: number | null;
+            postId: number | null;
+            author: {
+              id: number;
+              name: string;
+              verified: boolean;
+              invitedBy: number | null;
+            } | null;
+          }[];
+        }[];
+      }[]
+    >();
+
+    ctx.expect(response.length).eq(1);
+    ctx.expect(response[0]?.posts.length).eq(1);
+
+    ctx.expect(response[0]?.posts[0]?.comments.length).eq(1);
+
+    ctx.expect(response[0]).toEqual({
+      id: 1,
+      name: 'Dan',
+      verified: false,
+      invitedBy: null,
+      posts: [
+        {
+          id: 1,
+          ownerId: 1,
+          content: 'Post1',
+          createdAt: response[0]?.posts[0]?.createdAt,
+          comments: [
+            {
+              id: 1,
+              content: 'Comment1',
+              creator: 2,
+              author: {
+                id: 2,
+                name: 'Andrew',
+                verified: false,
+                invitedBy: null
+              },
+              postId: 1,
+              createdAt: response[0]?.posts[0]?.comments[0]?.createdAt
+            }
+          ]
+        }
+      ]
+    });
+  });
+
   /*
-	One three-level relation + 1 first-level relatioon
+	One three-level relation + 1 first-level relation
 	1. users+posts+comments+comment_owner
 	2. users+users
 */
 
   /*
-	One four-level relation users+posts+comments+coment_likes
+	One four-level relation users+posts+comments+comment_likes
 */
 
   /*
@@ -4158,9 +4147,11 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
           columns: {},
           with: {
             group: true
-          }
+          },
+          orderBy: usersToGroupsTable.userId
         }
-      }
+      },
+      orderBy: usersTable.id
     });
 
     expectTypeOf(response).toEqualTypeOf<
@@ -4224,22 +4215,22 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
       name: 'Alex',
       verified: false,
       invitedBy: null,
-      usersToGroups: [
+      usersToGroups: ctx.expect.arrayContaining([
         {
           group: {
-            id: 3,
-            name: 'Group3',
+            id: 2,
+            name: 'Group2',
             description: null
           }
         },
         {
           group: {
-            id: 2,
-            name: 'Group2',
+            id: 3,
+            name: 'Group3',
             description: null
           }
         }
-      ]
+      ])
     });
   });
 
@@ -6035,7 +6026,8 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
                 lower: sql<string>`lower(${groupsTable.name})`.as('lower_name')
               }
             }
-          }
+          },
+          orderBy: usersToGroupsTable.groupId
         }
       }
     });
@@ -6111,17 +6103,17 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
       usersToGroups: [
         {
           group: {
-            id: 3,
-            name: 'Group3',
-            lower: 'group3',
+            id: 2,
+            name: 'Group2',
+            lower: 'group2',
             description: null
           }
         },
         {
           group: {
-            id: 2,
-            name: 'Group2',
-            lower: 'group2',
+            id: 3,
+            name: 'Group3',
+            lower: 'group3',
             description: null
           }
         }
@@ -6257,6 +6249,23 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
     });
   });
 
+  test('Filter by columns not present in select', async (ctx) => {
+    await ctx.db.insert(usersTable).values([
+      { id: 1, name: 'Dan' },
+      { id: 2, name: 'Andrew' },
+      { id: 3, name: 'Alex' }
+    ]);
+
+    const response = await ctx.db.query.usersTable.findFirst({
+      columns: {
+        id: true
+      },
+      where: eq(usersTable.name, 'Dan')
+    });
+
+    ctx.expect(response).toEqual({ id: 1 });
+  });
+
   test('.toSQL()', (ctx) => {
     const query = ctx.db.query.usersTable.findFirst().toSQL();
 
@@ -6277,12 +6286,12 @@ describe.concurrent.each([{ type: 'pg' } /**{ type: 'http' }**/])('Drizzle $type
   });
 });
 
-async function waitForReplication(): Promise<void> {
+async function waitForReplication(dbName: string): Promise<void> {
   try {
-    await api.branches.getBranchList({ workspace, database, region });
-  } catch (error) {
-    console.log(`Waiting for create database replication to finish...`);
     await new Promise((resolve) => setTimeout(resolve, 2000));
-    return await waitForReplication();
+    await api.branch.getBranchList({ pathParams: { workspace, dbName, region } });
+  } catch (error) {
+    console.log(`Replication not ready yet, retrying...`);
+    return await waitForReplication(dbName);
   }
 }