Merge pull request #1578 from L-Mario564/extend-with-clause

Improve `with` clause
drizzle-team · Dec 27, 2023 · 6d706c1 · 6d706c1
2 parents d5c45a6 + 846df6f
commit 6d706c1
Show file tree

Hide file tree

Showing 18 changed files with 995 additions and 69 deletions.
diff --git a/drizzle-orm/src/mysql-core/db.ts b/drizzle-orm/src/mysql-core/db.ts
@@ -147,6 +147,42 @@ export class MySqlDatabase<
 	with(...queries: WithSubquery[]) {
 		const self = this;
 
+		/**
+		 * Creates a select query.
+		 * 
+		 * Calling this method with no arguments will select all columns from the table. Pass a selection object to specify the columns you want to select.
+		 * 
+		 * Use `.from()` method to specify which table to select from.
+		 * 
+		 * See docs: {@link https://orm.drizzle.team/docs/select}
+		 * 
+		 * @param fields The selection object.
+		 * 
+		 * @example
+		 * 
+		 * ```ts
+		 * // Select all columns and all rows from the 'cars' table
+		 * const allCars: Car[] = await db.select().from(cars);
+		 * 
+		 * // Select specific columns and all rows from the 'cars' table
+		 * const carsIdsAndBrands: { id: number; brand: string }[] = await db.select({ 
+		 *   id: cars.id, 
+		 *   brand: cars.brand 
+		 * })
+		 *   .from(cars);
+		 * ```
+		 * 
+		 * Like in SQL, you can use arbitrary expressions as selection fields, not just table columns:
+		 * 
+		 * ```ts
+		 * // Select specific columns along with expression and all rows from the 'cars' table
+		 * const carsIdsAndLowerNames: { id: number; lowerBrand: string }[] = await db.select({
+		 *   id: cars.id,
+		 *   lowerBrand: sql<string>`lower(${cars.brand})`,
+		 * })
+		 *   .from(cars);
+		 * ```
+		 */
 		function select(): MySqlSelectBuilder<undefined, TPreparedQueryHKT>;
 		function select<TSelection extends SelectedFields>(
 			fields: TSelection,
@@ -160,6 +196,30 @@ export class MySqlDatabase<
 			});
 		}
 
+		/**
+		 * Adds `distinct` expression to the select query.
+		 * 
+		 * Calling this method will return only unique values. When multiple columns are selected, it returns rows with unique combinations of values in these columns.
+		 * 
+		 * Use `.from()` method to specify which table to select from.
+		 * 
+		 * See docs: {@link https://orm.drizzle.team/docs/select#distinct}
+		 * 
+		 * @param fields The selection object.
+		 * 
+		 * @example
+		 * ```ts
+		 * // Select all unique rows from the 'cars' table
+		 * await db.selectDistinct()
+		 *   .from(cars)
+		 *   .orderBy(cars.id, cars.brand, cars.color);
+		 * 
+		 * // Select all unique brands from the 'cars' table
+		 * await db.selectDistinct({ brand: cars.brand })
+		 *   .from(cars)
+		 *   .orderBy(cars.brand);
+		 * ```
+		 */
 		function selectDistinct(): MySqlSelectBuilder<undefined, TPreparedQueryHKT>;
 		function selectDistinct<TSelection extends SelectedFields>(
 			fields: TSelection,
@@ -176,7 +236,55 @@ export class MySqlDatabase<
 			});
 		}
 
-		return { select, selectDistinct };
+		/**
+		 * Creates an update query.
+		 * 
+		 * Calling this method without `.where()` clause will update all rows in a table. The `.where()` clause specifies which rows should be updated.
+		 * 
+		 * Use `.set()` method to specify which values to update.
+		 * 
+		 * See docs: {@link https://orm.drizzle.team/docs/update} 
+		 * 
+		 * @param table The table to update.
+		 * 
+		 * @example
+		 * 
+		 * ```ts
+		 * // Update all rows in the 'cars' table
+		 * await db.update(cars).set({ color: 'red' });
+		 * 
+		 * // Update rows with filters and conditions
+		 * await db.update(cars).set({ color: 'red' }).where(eq(cars.brand, 'BMW'));
+		 * ```
+		 */
+		function update<TTable extends MySqlTable>(table: TTable): MySqlUpdateBuilder<TTable, TQueryResult, TPreparedQueryHKT> {
+			return new MySqlUpdateBuilder(table, self.session, self.dialect, queries);
+		}
+
+		/**
+		 * Creates a delete query.
+		 * 
+		 * Calling this method without `.where()` clause will delete all rows in a table. The `.where()` clause specifies which rows should be deleted. 
+		 * 
+		 * See docs: {@link https://orm.drizzle.team/docs/delete}
+		 *  
+		 * @param table The table to delete from.
+		 * 
+		 * @example
+		 * 
+		 * ```ts
+		 * // Delete all rows in the 'cars' table
+		 * await db.delete(cars);
+		 * 
+		 * // Delete rows with filters and conditions
+		 * await db.delete(cars).where(eq(cars.color, 'green'));
+		 * ```
+		 */
+		function delete_<TTable extends MySqlTable>(table: TTable): MySqlDeleteBase<TTable, TQueryResult, TPreparedQueryHKT> {
+			return new MySqlDeleteBase(table, self.session, self.dialect, queries);
+		}
+
+		return { select, selectDistinct, update, delete: delete_ };
 	}
 
 	/**

diff --git a/drizzle-orm/src/mysql-core/dialect.ts b/drizzle-orm/src/mysql-core/dialect.ts
@@ -79,14 +79,30 @@ export class MySqlDialect {
 		return `'${str.replace(/'/g, "''")}'`;
 	}
 
-	buildDeleteQuery({ table, where, returning }: MySqlDeleteConfig): SQL {
+	private buildWithCTE(queries: Subquery[] | undefined): SQL | undefined {
+		if (!queries?.length) return undefined;
+
+		const withSqlChunks = [sql`with `];
+		for (const [i, w] of queries.entries()) {
+			withSqlChunks.push(sql`${sql.identifier(w[SubqueryConfig].alias)} as (${w[SubqueryConfig].sql})`);
+			if (i < queries.length - 1) {
+				withSqlChunks.push(sql`, `);
+			}
+		}
+		withSqlChunks.push(sql` `);
+		return sql.join(withSqlChunks);
+	}
+
+	buildDeleteQuery({ table, where, returning, withList }: MySqlDeleteConfig): SQL {
+		const withSql = this.buildWithCTE(withList);
+
 		const returningSql = returning
 			? sql` returning ${this.buildSelection(returning, { isSingleTable: true })}`
 			: undefined;
 
 		const whereSql = where ? sql` where ${where}` : undefined;
 
-		return sql`delete from ${table}${whereSql}${returningSql}`;
+		return sql`${withSql}delete from ${table}${whereSql}${returningSql}`;
 	}
 
 	buildUpdateSet(table: MySqlTable, set: UpdateSet): SQL {
@@ -106,7 +122,9 @@ export class MySqlDialect {
 		);
 	}
 
-	buildUpdateQuery({ table, set, where, returning }: MySqlUpdateConfig): SQL {
+	buildUpdateQuery({ table, set, where, returning, withList }: MySqlUpdateConfig): SQL {
+		const withSql = this.buildWithCTE(withList);
+
 		const setSql = this.buildUpdateSet(table, set);
 
 		const returningSql = returning
@@ -115,7 +133,7 @@ export class MySqlDialect {
 
 		const whereSql = where ? sql` where ${where}` : undefined;
 
-		return sql`update ${table} set ${setSql}${whereSql}${returningSql}`;
+		return sql`${withSql}update ${table} set ${setSql}${whereSql}${returningSql}`;
 	}
 
 	/**
@@ -226,18 +244,7 @@ export class MySqlDialect {
 
 		const isSingleTable = !joins || joins.length === 0;
 
-		let withSql: SQL | undefined;
-		if (withList?.length) {
-			const withSqlChunks = [sql`with `];
-			for (const [i, w] of withList.entries()) {
-				withSqlChunks.push(sql`${sql.identifier(w[SubqueryConfig].alias)} as (${w[SubqueryConfig].sql})`);
-				if (i < withList.length - 1) {
-					withSqlChunks.push(sql`, `);
-				}
-			}
-			withSqlChunks.push(sql` `);
-			withSql = sql.join(withSqlChunks);
-		}
+		const withSql = this.buildWithCTE(withList);
 
 		const distinctSql = distinct ? sql` distinct` : undefined;
 

diff --git a/drizzle-orm/src/mysql-core/query-builders/delete.ts b/drizzle-orm/src/mysql-core/query-builders/delete.ts
@@ -13,6 +13,7 @@ import type { MySqlTable } from '~/mysql-core/table.ts';
 import { QueryPromise } from '~/query-promise.ts';
 import type { Query, SQL, SQLWrapper } from '~/sql/sql.ts';
 import type { SelectedFieldsOrdered } from './select.types.ts';
+import type { Subquery } from '~/subquery.ts';
 
 export type MySqlDeleteWithout<
 	T extends AnyMySqlDeleteBase,
@@ -40,6 +41,7 @@ export interface MySqlDeleteConfig {
 	where?: SQL | undefined;
 	table: MySqlTable;
 	returning?: SelectedFieldsOrdered;
+	withList?: Subquery[];
 }
 
 export type MySqlDeletePrepare<T extends AnyMySqlDeleteBase> = PreparedQueryKind<
@@ -92,9 +94,10 @@ export class MySqlDeleteBase<
 		private table: TTable,
 		private session: MySqlSession,
 		private dialect: MySqlDialect,
+		withList?: Subquery[],
 	) {
 		super();
-		this.config = { table };
+		this.config = { table, withList };
 	}
 
 	/** 

diff --git a/drizzle-orm/src/mysql-core/query-builders/update.ts b/drizzle-orm/src/mysql-core/query-builders/update.ts
@@ -15,12 +15,14 @@ import { QueryPromise } from '~/query-promise.ts';
 import type { Query, SQL, SQLWrapper } from '~/sql/sql.ts';
 import { mapUpdateSet, type UpdateSet } from '~/utils.ts';
 import type { SelectedFieldsOrdered } from './select.types.ts';
+import type { Subquery } from '~/subquery.ts';
 
 export interface MySqlUpdateConfig {
 	where?: SQL | undefined;
 	set: UpdateSet;
 	table: MySqlTable;
 	returning?: SelectedFieldsOrdered;
+	withList?: Subquery[];
 }
 
 export type MySqlUpdateSetSource<TTable extends MySqlTable> =
@@ -46,10 +48,11 @@ export class MySqlUpdateBuilder<
 		private table: TTable,
 		private session: MySqlSession,
 		private dialect: MySqlDialect,
+		private withList?: Subquery[],
 	) {}
 
 	set(values: MySqlUpdateSetSource<TTable>): MySqlUpdateBase<TTable, TQueryResult, TPreparedQueryHKT> {
-		return new MySqlUpdateBase(this.table, mapUpdateSet(this.table, values), this.session, this.dialect);
+		return new MySqlUpdateBase(this.table, mapUpdateSet(this.table, values), this.session, this.dialect, this.withList);
 	}
 }
 
@@ -126,9 +129,10 @@ export class MySqlUpdateBase<
 		set: UpdateSet,
 		private session: MySqlSession,
 		private dialect: MySqlDialect,
+		withList?: Subquery[],
 	) {
 		super();
-		this.config = { set, table };
+		this.config = { set, table, withList };
 	}
 
 	/**