add IDENTITY column support. (kysely-org#823)

thecodrr · Sep 3, 2024 · 9bf7c27 · 9bf7c27
1 parent d11c8ae
commit 9bf7c27
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 16 deletions.
diff --git a/src/operation-node/column-definition-node.ts b/src/operation-node/column-definition-node.ts
@@ -4,7 +4,6 @@ import { ColumnNode } from './column-node.js'
 import { DefaultValueNode } from './default-value-node.js'
 import { GeneratedNode } from './generated-node.js'
 import { OperationNode } from './operation-node.js'
-import { RawNode } from './raw-node.js'
 import { ReferencesNode } from './references-node.js'
 
 export type ColumnDefinitionNodeProps = Omit<
@@ -28,6 +27,7 @@ export interface ColumnDefinitionNode extends OperationNode {
   readonly frontModifiers?: ReadonlyArray<OperationNode>
   readonly endModifiers?: ReadonlyArray<OperationNode>
   readonly nullsNotDistinct?: boolean
+  readonly identity?: boolean
 }
 
 /**

diff --git a/src/operation-node/operation-node-transformer.ts b/src/operation-node/operation-node-transformer.ts
@@ -485,6 +485,7 @@ export class OperationNodeTransformer {
       frontModifiers: this.transformNodeList(node.frontModifiers),
       endModifiers: this.transformNodeList(node.endModifiers),
       nullsNotDistinct: node.nullsNotDistinct,
+      identity: node.identity,
     })
   }
 

diff --git a/src/query-compiler/default-query-compiler.ts b/src/query-compiler/default-query-compiler.ts
@@ -684,6 +684,10 @@ export class DefaultQueryCompiler
       this.visitNode(node.generated)
     }
 
+    if (node.identity) {
+      this.append(' identity')
+    }
+
     if (node.defaultTo) {
       this.append(' ')
       this.visitNode(node.defaultTo)

diff --git a/src/schema/column-definition-builder.ts b/src/schema/column-definition-builder.ts
@@ -37,6 +37,19 @@ export class ColumnDefinitionBuilder implements OperationNodeSource {
     )
   }
 
+  /**
+   * Makes the column an identity column.
+   *
+   * This only works on some dialects like MS SQL Server (MSSQL).
+   *
+   * For PostgreSQL's `generated always as identity` use {@link generatedAlwaysAsIdentity}.
+   */
+  identity(): ColumnDefinitionBuilder {
+    return new ColumnDefinitionBuilder(
+      ColumnDefinitionNode.cloneWith(this.#node, { identity: true })
+    )
+  }
+
   /**
    * Makes the column the primary key.
    *
@@ -248,7 +261,11 @@ export class ColumnDefinitionBuilder implements OperationNodeSource {
   }
 
   /**
-   * Adds the `generated always as identity` specifier on supported dialects.
+   * Adds the `generated always as identity` specifier.
+   *
+   * This only works on some dialects like PostgreSQL.
+   *
+   * For MS SQL Server (MSSQL)'s identity column use {@link identity}.
    */
   generatedAlwaysAsIdentity(): ColumnDefinitionBuilder {
     return new ColumnDefinitionBuilder(

diff --git a/test/node/src/introspect.test.ts b/test/node/src/introspect.test.ts
@@ -831,10 +831,7 @@ for (const dialect of DIALECTS) {
         await ctx.db.schema
           .createTable('some_schema.pet')
           .addColumn('some_column', 'integer', (col) =>
-            col
-              .notNull()
-              .modifyFront(sql`identity(1,1)`)
-              .primaryKey(),
+            col.identity().notNull().primaryKey()
           )
           .addColumn('some_column_plus_1', sql``, (col) =>
             col.modifyEnd(sql`as (some_column + 1)`),

diff --git a/test/node/src/schema.test.ts b/test/node/src/schema.test.ts
@@ -299,10 +299,7 @@ for (const dialect of DIALECTS) {
           const builder = ctx.db.schema
             .createTable('test')
             .addColumn('a', 'integer', (col) =>
-              col
-                .notNull()
-                .modifyFront(sql`identity(1,1)`)
-                .primaryKey(),
+              col.identity().notNull().primaryKey()
             )
             .addColumn('b', 'integer', (col) =>
               col
@@ -343,7 +340,7 @@ for (const dialect of DIALECTS) {
             mssql: {
               sql: [
                 'create table "test"',
-                '("a" integer identity(1,1) not null primary key,',
+                '("a" integer identity not null primary key,',
                 '"b" integer references "test" ("a") on delete no action on update no action check (b < 10),',
                 '"c" varchar,',
                 '"d" varchar(10),',

diff --git a/test/node/src/test-setup.ts b/test/node/src/test-setup.ts
@@ -373,11 +373,7 @@ export function createTableWithId(
 
   if (dialect === 'mssql') {
     return builder.addColumn('id', 'integer', (col) =>
-      col
-        .notNull()
-        // TODO: change to method when its implemented
-        .modifyFront(sql`identity(1,1)`)
-        .primaryKey(),
+      col.identity().notNull().primaryKey()
     )
   }