feat: support IF EXISTS in DROP COLUMN (#1692)

Co-authored-by: Igal Klebanov <igalklebanov@gmail.com>

Author: shuaixr

src/dialect/mssql/mssql-query-compiler.ts

@@ -52,25 +52,38 @@ export class MssqlQueryCompiler extends DefaultQueryCompiler {
     // multiple of these are not really supported by mssql,
     // but for the sake of WYSIWYG.
     if (nodesByKind.AlterColumnNode) {
-      if (!first) this.append(', ')
+      if (!first) {
+        this.append(', ')
+      }
+
       this.compileList(nodesByKind.AlterColumnNode)
     }
 
     if (nodesByKind.DropColumnNode) {
-      if (!first) this.append(', ')
-      this.append('drop column ')
+      if (!first) {
+        this.append(', ')
+      }
+
+      this.append('drop ')
       this.compileList(nodesByKind.DropColumnNode)
+      first = false
     }
 
     // not really supported by mssql, but for the sake of WYSIWYG.
     if (nodesByKind.ModifyColumnNode) {
-      if (!first) this.append(', ')
+      if (!first) {
+        this.append(', ')
+      }
+
       this.compileList(nodesByKind.ModifyColumnNode)
     }
 
     // not really supported by mssql, but for the sake of WYSIWYG.
     if (nodesByKind.RenameColumnNode) {
-      if (!first) this.append(', ')
+      if (!first) {
+        this.append(', ')
+      }
+
       this.compileList(nodesByKind.RenameColumnNode)
     }
   }
@@ -80,6 +93,12 @@ export class MssqlQueryCompiler extends DefaultQueryCompiler {
   }
 
   protected override visitDropColumn(node: DropColumnNode): void {
+    this.append('column ')
+
+    if (node.ifExists) {
+      this.append('if exists ')
+    }
+
     this.visitNode(node.column)
   }
 

src/index.ts

@@ -66,6 +66,7 @@ export * from './schema/create-view-builder.js'
 export * from './schema/refresh-materialized-view-builder.js'
 export * from './schema/drop-view-builder.js'
 export * from './schema/alter-column-builder.js'
+export * from './schema/drop-column-builder.js'
 
 export * from './dynamic/dynamic.js'
 export * from './dynamic/dynamic-reference-builder.js'

src/operation-node/drop-column-node.ts

@@ -2,14 +2,21 @@ import type { OperationNode } from './operation-node.js'
 import { freeze } from '../util/object-utils.js'
 import { ColumnNode } from './column-node.js'
 
+export type DropColumnNodeProps = Omit<DropColumnNode, 'kind' | 'column'>
+
 export interface DropColumnNode extends OperationNode {
   readonly kind: 'DropColumnNode'
   readonly column: ColumnNode
+  readonly ifExists?: boolean
 }
 
 type DropColumnNodeFactory = Readonly<{
   is(node: OperationNode): node is DropColumnNode
   create(column: string): Readonly<DropColumnNode>
+  cloneWith(
+    node: DropColumnNode,
+    props: DropColumnNodeProps,
+  ): Readonly<DropColumnNode>
 }>
 
 /**
@@ -27,4 +34,11 @@ export const DropColumnNode: DropColumnNodeFactory =
         column: ColumnNode.create(column),
       })
     },
+
+    cloneWith(node, props) {
+      return freeze({
+        ...node,
+        ...props,
+      })
+    },
   })

src/operation-node/operation-node-transformer.ts

@@ -865,6 +865,7 @@ export class OperationNodeTransformer {
     return {
       kind: 'DropColumnNode',
       column: this.transformNode(node.column, queryId),
+      ifExists: node.ifExists,
     } satisfies AllProps<DropColumnNode>
   }
 

src/query-compiler/default-query-compiler.ts

@@ -1232,9 +1232,13 @@ export class DefaultQueryCompiler
 
   protected override visitDropColumn(node: DropColumnNode): void {
     this.append('drop column ')
+
+    if (node.ifExists) {
+      this.append('if exists ')
+    }
+
     this.visitNode(node.column)
   }
-
   protected override visitAlterColumn(node: AlterColumnNode): void {
     this.append('alter column ')
     this.visitNode(node.column)

src/schema/alter-table-builder.ts

@@ -59,6 +59,10 @@ import {
   type ExpressionOrFactory,
   parseExpression,
 } from '../parser/expression-parser.js'
+import {
+  DropColumnBuilder,
+  type DropColumnBuilderCallback,
+} from './drop-column-builder.js'
 
 /**
  * This builder can be used to create a `alter table` query.
@@ -103,12 +107,19 @@ export class AlterTableBuilder implements ColumnAlteringInterface {
     })
   }
 
-  dropColumn(column: string): AlterTableColumnAlteringBuilder {
+  dropColumn(
+    column: string,
+    build: DropColumnBuilderCallback = noop,
+  ): AlterTableColumnAlteringBuilder {
+    const builder = build(
+      new DropColumnBuilder({ node: DropColumnNode.create(column) }),
+    )
+
     return new AlterTableColumnAlteringBuilder({
       ...this.#props,
       node: AlterTableNode.cloneWithColumnAlteration(
         this.#props.node,
-        DropColumnNode.create(column),
+        builder.toOperationNode(),
       ),
     })
   }
@@ -431,12 +442,19 @@ export class AlterTableColumnAlteringBuilder
     })
   }
 
-  dropColumn(column: string): AlterTableColumnAlteringBuilder {
+  dropColumn(
+    column: string,
+    build: DropColumnBuilderCallback = noop,
+  ): AlterTableColumnAlteringBuilder {
+    const builder = build(
+      new DropColumnBuilder({ node: DropColumnNode.create(column) }),
+    )
+
     return new AlterTableColumnAlteringBuilder({
       ...this.#props,
       node: AlterTableNode.cloneWithColumnAlteration(
         this.#props.node,
-        DropColumnNode.create(column),
+        builder.toOperationNode(),
       ),
     })
   }

src/schema/drop-column-builder.ts

@@ -0,0 +1,30 @@
+import { DropColumnNode } from '../operation-node/drop-column-node.js'
+import type { OperationNodeSource } from '../operation-node/operation-node-source.js'
+import { freeze } from '../util/object-utils.js'
+
+export class DropColumnBuilder implements OperationNodeSource {
+  readonly #props: DropColumnBuilderProps
+
+  constructor(props: DropColumnBuilderProps) {
+    this.#props = freeze({ ...props })
+  }
+
+  ifExists(): DropColumnBuilder {
+    return new DropColumnBuilder({
+      ...this.#props,
+      node: DropColumnNode.cloneWith(this.#props.node, { ifExists: true }),
+    })
+  }
+
+  toOperationNode(): DropColumnNode {
+    return this.#props.node
+  }
+}
+
+export interface DropColumnBuilderProps {
+  readonly node: DropColumnNode
+}
+
+export type DropColumnBuilderCallback = (
+  builder: DropColumnBuilder,
+) => DropColumnBuilder

test/node/src/schema.test.ts

@@ -2954,6 +2954,7 @@ for (const dialect of DIALECTS) {
           .createTable('test')
           .addColumn('varchar_col', 'varchar(255)')
           .addColumn('integer_col', 'integer')
+          .addColumn('different_col', 'text')
           .execute()
       })
 
@@ -3397,6 +3398,29 @@ for (const dialect of DIALECTS) {
           await builder.execute()
         })
 
+        if (sqlSpec === 'postgres' || sqlSpec === 'mssql') {
+          it('should drop a column if it exists', async () => {
+            const builder = ctx.db.schema
+              .alterTable('test')
+              .dropColumn('varchar_col', (col) => col.ifExists())
+
+            testSql(builder, dialect, {
+              postgres: {
+                sql: 'alter table "test" drop column if exists "varchar_col"',
+                parameters: [],
+              },
+              mysql: NOT_SUPPORTED,
+              mssql: {
+                sql: 'alter table "test" drop column if exists "varchar_col"',
+                parameters: [],
+              },
+              sqlite: NOT_SUPPORTED,
+            })
+
+            await builder.execute()
+          })
+        }
+
         if (
           sqlSpec === 'postgres' ||
           sqlSpec === 'mysql' ||
@@ -3434,7 +3458,7 @@ for (const dialect of DIALECTS) {
                 sql: [
                   'alter table "test"',
                   'drop column "varchar_col",',
-                  '"text_col"',
+                  'column "text_col"',
                 ],
                 parameters: [],
               },
@@ -3453,6 +3477,41 @@ for (const dialect of DIALECTS) {
         }
       })
 
+      if (sqlSpec === 'postgres' || sqlSpec === 'mssql') {
+        it('should drop multiple columns if it exists', async () => {
+          const builder = ctx.db.schema
+            .alterTable('test')
+            .dropColumn('varchar_col')
+            .dropColumn('not_exist_col', (col) => col.ifExists())
+            .dropColumn('different_col')
+
+          testSql(builder, dialect, {
+            postgres: {
+              sql: [
+                'alter table "test"',
+                'drop column "varchar_col",',
+                'drop column if exists "not_exist_col",',
+                'drop column "different_col"',
+              ],
+              parameters: [],
+            },
+            mysql: NOT_SUPPORTED,
+            mssql: {
+              sql: [
+                'alter table "test"',
+                'drop column "varchar_col",',
+                'column if exists "not_exist_col",',
+                'column "different_col"',
+              ],
+              parameters: [],
+            },
+            sqlite: NOT_SUPPORTED,
+          })
+
+          await builder.execute()
+        })
+      }
+
       if (
         sqlSpec === 'postgres' ||
         sqlSpec === 'mysql' ||