Add with clause tests

Add tests for `db.with().insert()`, `db.with().delete()` and `db.with().update()` for all dialects
L-Mario564 · Nov 29, 2023 · 217369a · 217369a
1 parent 68b351f
commit 217369a
Show file tree

Hide file tree

Showing 3 changed files with 372 additions and 1 deletion.
diff --git a/integration-tests/tests/libsql.test.ts b/integration-tests/tests/libsql.test.ts
@@ -22,7 +22,8 @@ import {
 	sum,
 	sumDistinct,
 	max,
-	min
+	min,
+	lt
 } from 'drizzle-orm';
 import { drizzle, type LibSQLDatabase } from 'drizzle-orm/libsql';
 import { migrate } from 'drizzle-orm/libsql/migrator';
@@ -36,6 +37,7 @@ import {
 	int,
 	integer,
 	intersect,
+	numeric,
 	primaryKey,
 	sqliteTable,
 	sqliteTableCreator,
@@ -1243,6 +1245,131 @@ test.serial('with ... select', async (t) => {
 	]);
 });
 
+test.serial('with ... update', async (t) => {
+	const { db } = t.context;
+
+	const products = sqliteTable('products', {
+		id: integer('id').primaryKey(),
+		price: numeric('price').notNull(),
+		cheap: integer('cheap', { mode: 'boolean' }).notNull().default(false)
+	});
+
+	await db.run(sql`drop table if exists ${products}`);
+	await db.run(sql`create table ${products} (
+		id integer primary key,
+		price numeric not null,
+		cheap integer not null default 0
+	)`);
+
+	await db.insert(products).values([
+		{ price: '10.99' },
+		{ price: '25.85' },
+		{ price: '32.99' },
+		{ price: '2.50' },
+		{ price: '4.59' },
+	]);
+
+	const averagePrice = db
+		.$with('average_price')
+		.as(
+			db
+				.select({
+					value: sql`avg(${products.price})`.as('value')
+				})
+				.from(products)
+		);
+
+	const result = await db
+		.with(averagePrice)
+		.update(products)
+		.set({
+			cheap: true
+		})
+		.where(lt(products.price, sql`(select * from ${averagePrice})`))
+		.returning({
+			id: products.id
+		});
+
+	t.deepEqual(result, [
+		{ id: 1 },
+		{ id: 4 },
+		{ id: 5 }
+	]);
+});
+
+test.serial('with ... insert', async (t) => {
+	const { db } = t.context;
+
+	const users = sqliteTable('users', {
+		username: text('username').notNull(),
+		admin: integer('admin', { mode: 'boolean' }).notNull()
+	});
+
+	await db.run(sql`drop table if exists ${users}`);
+	await db.run(sql`create table ${users} (username text not null, admin integer not null default 0)`);
+
+	const userCount = db
+		.$with('user_count')
+		.as(
+			db
+				.select({
+					value: sql`count(*)`.as('value')
+				})
+				.from(users)
+		);
+
+	const result = await db
+		.with(userCount)
+		.insert(users)
+		.values([
+			{ username: 'user1', admin: sql`((select * from ${userCount}) = 0)` }
+		])
+		.returning({
+			admin: users.admin
+		});
+
+	t.deepEqual(result, [{ admin: true }])
+});
+
+test.serial('with ... delete', async (t) => {
+	const { db } = t.context;
+
+	await db.insert(orders).values([
+		{ region: 'Europe', product: 'A', amount: 10, quantity: 1 },
+		{ region: 'Europe', product: 'A', amount: 20, quantity: 2 },
+		{ region: 'Europe', product: 'B', amount: 20, quantity: 2 },
+		{ region: 'Europe', product: 'B', amount: 30, quantity: 3 },
+		{ region: 'US', product: 'A', amount: 30, quantity: 3 },
+		{ region: 'US', product: 'A', amount: 40, quantity: 4 },
+		{ region: 'US', product: 'B', amount: 40, quantity: 4 },
+		{ region: 'US', product: 'B', amount: 50, quantity: 5 },
+	]);
+
+	const averageAmount = db
+		.$with('average_amount')
+		.as(
+			db
+				.select({
+					value: sql`avg(${orders.amount})`.as('value')
+				})
+				.from(orders)
+		);
+
+	const result = await db
+		.with(averageAmount)
+		.delete(orders)
+		.where(gt(orders.amount, sql`(select * from ${averageAmount})`))
+		.returning({
+			id: orders.id
+		});
+
+	t.deepEqual(result, [
+		{ id: 6 },
+		{ id: 7 },
+		{ id: 8 }
+	]);
+});
+
 test.serial('select from subquery sql', async (t) => {
 	const { db } = t.context;
 

diff --git a/integration-tests/tests/mysql.test.ts b/integration-tests/tests/mysql.test.ts
@@ -23,13 +23,15 @@ import {
 	avgDistinct,
 	max,
 	min,
+	lt,
 } from 'drizzle-orm';
 import {
 	alias,
 	bigint,
 	boolean,
 	date,
 	datetime,
+	decimal,
 	except,
 	exceptAll,
 	foreignKey,
@@ -1483,6 +1485,122 @@ test.serial('with ... select', async (t) => {
 	]);
 });
 
+test.serial('with ... update', async (t) => {
+	const { db } = t.context;
+
+	const products = mysqlTable('products', {
+		id: serial('id').primaryKey(),
+		price: decimal('price', {
+			precision: 15,
+			scale: 2
+		}).notNull(),
+		cheap: boolean('cheap').notNull().default(false)
+	});
+
+	await db.execute(sql`drop table if exists ${products}`);
+	await db.execute(sql`create table ${products} (
+		id serial primary key,
+		price decimal(15, 2) not null,
+		cheap boolean not null default false
+	)`);
+
+	await db.insert(products).values([
+		{ price: '10.99' },
+		{ price: '25.85' },
+		{ price: '32.99' },
+		{ price: '2.50' },
+		{ price: '4.59' },
+	]);
+
+	const averagePrice = db
+		.$with('average_price')
+		.as(
+			db
+				.select({
+					value: sql`avg(${products.price})`.as('value')
+				})
+				.from(products)
+		);
+
+	await db
+		.with(averagePrice)
+		.update(products)
+		.set({
+			cheap: true
+		})
+		.where(lt(products.price, sql`(select * from ${averagePrice})`));
+
+	const result = await db
+		.select({
+			id: products.id
+		})
+		.from(products)
+		.where(eq(products.cheap, true))
+
+	t.deepEqual(result, [
+		{ id: 1 },
+		{ id: 4 },
+		{ id: 5 }
+	]);
+});
+
+test.serial('with ... delete', async (t) => {
+	const { db } = t.context;
+
+	await db.execute(sql`drop table if exists \`orders\``);
+	await db.execute(
+		sql`
+			create table \`orders\` (
+				\`id\` serial primary key,
+				\`region\` text not null,
+				\`product\` text not null,
+				\`amount\` int not null,
+				\`quantity\` int not null
+			)
+		`,
+	);
+
+	await db.insert(orders).values([
+		{ region: 'Europe', product: 'A', amount: 10, quantity: 1 },
+		{ region: 'Europe', product: 'A', amount: 20, quantity: 2 },
+		{ region: 'Europe', product: 'B', amount: 20, quantity: 2 },
+		{ region: 'Europe', product: 'B', amount: 30, quantity: 3 },
+		{ region: 'US', product: 'A', amount: 30, quantity: 3 },
+		{ region: 'US', product: 'A', amount: 40, quantity: 4 },
+		{ region: 'US', product: 'B', amount: 40, quantity: 4 },
+		{ region: 'US', product: 'B', amount: 50, quantity: 5 },
+	]);
+
+	const averageAmount = db
+		.$with('average_amount')
+		.as(
+			db
+				.select({
+					value: sql`avg(${orders.amount})`.as('value')
+				})
+				.from(orders)
+		);
+
+	await db
+		.with(averageAmount)
+		.delete(orders)
+		.where(gt(orders.amount, sql`(select * from ${averageAmount})`));
+
+	const result = await db
+		.select({
+			id: orders.id
+		})
+		.from(orders);
+
+	t.deepEqual(result, [
+		{ id: 1 },
+		{ id: 2 },
+		{ id: 3 },
+		{ id: 4 },
+		{ id: 5 }
+	]);
+});
+
 test.serial('select from subquery sql', async (t) => {
 	const { db } = t.context;