From b20a6ba89e7eb411d0409eb099fee30c3433f414 Mon Sep 17 00:00:00 2001
From: Igal Klebanov <igalklebanov@gmail.com>
Date: Wed, 14 Feb 2024 18:43:21 +0200
Subject: [PATCH] add `IDENTITY` column support. (#823)

---
 src/operation-node/column-definition-node.ts  |  2 +-
 .../operation-node-transformer.ts             |  1 +
 src/query-compiler/default-query-compiler.ts  |  4 ++++
 src/schema/column-definition-builder.ts       | 19 ++++++++++++++++++-
 test/node/src/introspect.test.ts              |  5 +----
 test/node/src/schema.test.ts                  |  7 ++-----
 test/node/src/test-setup.ts                   |  6 +-----
 7 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/src/operation-node/column-definition-node.ts b/src/operation-node/column-definition-node.ts
index bdd61c1e3..395d464b4 100644
--- 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
index 681be2057..ac8400755 100644
--- a/src/operation-node/operation-node-transformer.ts
+++ b/src/operation-node/operation-node-transformer.ts
@@ -443,6 +443,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
index a6776ac80..aea3e3530 100644
--- a/src/query-compiler/default-query-compiler.ts
+++ b/src/query-compiler/default-query-compiler.ts
@@ -603,6 +603,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
index 877295291..61d4961e1 100644
--- 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
index 7daf17cd9..86e1a6cac 100644
--- 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
index b21ae322e..936f50d46 100644
--- 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
index 00738c58b..f05753540 100644
--- 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()
     )
   }