Rework index columns op class handling

drizzle-team · May 15, 2024 · 3041c82 · 3041c82
1 parent a420907
commit 3041c82
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 24 deletions.
diff --git a/drizzle-orm/src/column-builder.ts b/drizzle-orm/src/column-builder.ts
@@ -1,7 +1,7 @@
 import { entityKind } from '~/entity.ts';
 import type { Column } from './column.ts';
 import type { MySqlColumn } from './mysql-core/index.ts';
-import type { IndexedColumn, PgColumn } from './pg-core/index.ts';
+import type { ExtraConfigColumn, PgColumn } from './pg-core/index.ts';
 import type { SQL } from './sql/sql.ts';
 import type { SQLiteColumn } from './sqlite-core/index.ts';
 import type { Simplify } from './utils.ts';
@@ -236,10 +236,8 @@ export type BuildColumn<
 	: never;
 
 export type BuildIndexColumn<
-	TTableName extends string,
-	TBuilder extends ColumnBuilderBase,
 	TDialect extends Dialect,
-> = TDialect extends 'pg' ? IndexedColumn<MakeColumnConfig<TBuilder['_'], TTableName>> : never;
+> = TDialect extends 'pg' ? ExtraConfigColumn : never;
 
 // TODO
 // try to make sql as well + indexRaw
@@ -259,12 +257,12 @@ export type BuildColumns<
 	& {};
 
 export type BuildExtraConfigColumns<
-	TTableName extends string,
+	_TTableName extends string,
 	TConfigMap extends Record<string, ColumnBuilderBase>,
 	TDialect extends Dialect,
 > =
 	& {
-		[Key in keyof TConfigMap]: BuildIndexColumn<TTableName, TConfigMap[Key], TDialect>;
+		[Key in keyof TConfigMap]: BuildIndexColumn<TDialect>;
 	}
 	& {};
 

diff --git a/drizzle-orm/src/pg-core/columns/common.ts b/drizzle-orm/src/pg-core/columns/common.ts
@@ -109,8 +109,8 @@ export abstract class PgColumnBuilder<
 	/** @internal */
 	buildExtraConfigColumn<TTableName extends string>(
 		table: AnyPgTable<{ name: TTableName }>,
-	): IndexedColumn {
-		return new IndexedColumn(table, this.config);
+	): ExtraConfigColumn {
+		return new ExtraConfigColumn(table, this.config);
 	}
 }
 
@@ -133,16 +133,27 @@ export abstract class PgColumn<
 	}
 }
 
-export class IndexedColumn<
+export type IndexedExtraConfigType = { order?: 'asc' | 'desc'; nulls?: 'first' | 'last'; opClass?: string };
+
+export class ExtraConfigColumn<
 	T extends ColumnBaseConfig<ColumnDataType, string> = ColumnBaseConfig<ColumnDataType, string>,
-> extends PgColumn<T, { order?: 'asc' | 'desc'; nulls?: 'first' | 'last'; opClass?: string }> {
-	static readonly [entityKind]: string = 'IndexColumn';
+> extends PgColumn<T, IndexedExtraConfigType> {
+	static readonly [entityKind]: string = 'ExtraConfigColumn';
 
 	override getSQLType(): string {
 		return this.getSQLType();
 	}
 
-	indexConfig = { order: this.config.order, nulls: this.config.nulls, op: this.config.opClass };
+	indexConfig: IndexedExtraConfigType = {
+		order: this.config.order ?? 'asc',
+		nulls: this.config.nulls ?? 'last',
+		opClass: this.config.opClass,
+	};
+	defaultConfig: IndexedExtraConfigType = {
+		order: 'asc',
+		nulls: 'last',
+		opClass: undefined,
+	};
 
 	asc(): Omit<this, 'asc' | 'desc'> {
 		this.indexConfig.order = 'asc';
@@ -165,11 +176,25 @@ export class IndexedColumn<
 	}
 
 	op(opClass: string): Omit<this, 'op'> {
-		this.indexConfig.op = opClass;
+		this.indexConfig.opClass = opClass;
 		return this;
 	}
 }
 
+export class IndexedColumn {
+	static readonly [entityKind]: string = 'IndexedColumn';
+	constructor(
+		name: string | undefined,
+		indexConfig: IndexedExtraConfigType,
+	) {
+		this.name = name;
+		this.indexConfig = indexConfig;
+	}
+
+	name: string | undefined;
+	indexConfig: IndexedExtraConfigType;
+}
+
 export type AnyPgColumn<TPartial extends Partial<ColumnBaseConfig<ColumnDataType, string>> = {}> = PgColumn<
 	Required<Update<ColumnBaseConfig<ColumnDataType, string>, TPartial>>
 >;

diff --git a/drizzle-orm/src/pg-core/indexes.ts b/drizzle-orm/src/pg-core/indexes.ts
@@ -1,7 +1,8 @@
-import type { SQL } from '~/sql/sql.ts';
+import { SQL } from '~/sql/sql.ts';
 
-import { entityKind } from '~/entity.ts';
-import type { IndexedColumn, PgColumn } from './columns/index.ts';
+import { entityKind, is } from '~/entity.ts';
+import type { ExtraConfigColumn, PgColumn } from './columns/index.ts';
+import { IndexedColumn } from './columns/index.ts';
 import type { PgTable } from './table.ts';
 
 interface IndexConfig {
@@ -47,16 +48,59 @@ export class IndexBuilderOn {
 
 	constructor(private unique: boolean, private name?: string) {}
 
-	on(...columns: [Partial<IndexedColumn> | SQL, ...Partial<IndexedColumn>[] | SQL[]]): IndexBuilder {
-		return new IndexBuilder(columns, this.unique, false, this.name);
+	on(...columns: [Partial<ExtraConfigColumn> | SQL, ...Partial<ExtraConfigColumn>[] | SQL[]]): IndexBuilder {
+		return new IndexBuilder(
+			columns.map((it) => {
+				if (is(it, SQL)) {
+					return it;
+				}
+				it = it as ExtraConfigColumn;
+				const clonedIndexedColumn = new IndexedColumn(it.name, it.indexConfig!);
+				it.indexConfig = JSON.parse(JSON.stringify(it.defaultConfig));
+				return clonedIndexedColumn;
+			}),
+			this.unique,
+			false,
+			this.name,
+		);
 	}
 
-	onOnly(...columns: [Partial<IndexedColumn | SQL>, ...Partial<IndexedColumn>[] | SQL[]]): IndexBuilder {
-		return new IndexBuilder(columns, this.unique, true, this.name);
+	onOnly(...columns: [Partial<ExtraConfigColumn | SQL>, ...Partial<ExtraConfigColumn>[] | SQL[]]): IndexBuilder {
+		return new IndexBuilder(
+			columns.map((it) => {
+				if (is(it, SQL)) {
+					return it;
+				}
+				it = it as ExtraConfigColumn;
+				const clonedIndexedColumn = new IndexedColumn(it.name, it.indexConfig!);
+				it.indexConfig = it.defaultConfig;
+				return clonedIndexedColumn;
+			}),
+			this.unique,
+			true,
+			this.name,
+		);
 	}
 
-	using(method: string, ...columns: [Partial<IndexedColumn | SQL>, ...Partial<IndexedColumn>[] | SQL[]]): IndexBuilder {
-		return new IndexBuilder(columns, this.unique, true, this.name, method);
+	using(
+		method: string,
+		...columns: [Partial<ExtraConfigColumn | SQL>, ...Partial<ExtraConfigColumn>[] | SQL[]]
+	): IndexBuilder {
+		return new IndexBuilder(
+			columns.map((it) => {
+				if (is(it, SQL)) {
+					return it;
+				}
+				it = it as ExtraConfigColumn;
+				const clonedIndexedColumn = new IndexedColumn(it.name, it.indexConfig!);
+				it.indexConfig = JSON.parse(JSON.stringify(it.defaultConfig));
+				return clonedIndexedColumn;
+			}),
+			this.unique,
+			true,
+			this.name,
+			method,
+		);
 	}
 }
 

diff --git a/drizzle-orm/src/table.ts b/drizzle-orm/src/table.ts
@@ -1,7 +1,7 @@
 import type { Column, GetColumnData } from './column.ts';
 import { entityKind } from './entity.ts';
 import type { OptionalKeyOnly, RequiredKeyOnly } from './operations.ts';
-import type { IndexedColumn } from './pg-core/index.ts';
+import type { ExtraConfigColumn } from './pg-core/index.ts';
 import type { SQLWrapper } from './sql/sql.ts';
 import type { Simplify, Update } from './utils.ts';
 
@@ -96,7 +96,7 @@ export class Table<T extends TableConfig = TableConfig> implements SQLWrapper {
 	[Columns]!: T['columns'];
 
 	/** @internal */
-	[ExtraConfigColumns]!: Record<string, IndexedColumn>;
+	[ExtraConfigColumns]!: Record<string, ExtraConfigColumn>;
 
 	/**
 	 *  @internal