Support dropping multiple types with schema.dropType(), cascade. (#1516)

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

Author: aantia

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

@@ -2,12 +2,17 @@ import { freeze } from '../util/object-utils.js'
 import { OperationNode } from './operation-node.js'
 import { SchemableIdentifierNode } from './schemable-identifier-node.js'
 
-export type DropTypeNodeParams = Omit<Partial<DropTypeNode>, 'kind' | 'name'>
+export type DropTypeNodeParams = Omit<
+  Partial<DropTypeNode>,
+  'kind' | 'name' | 'additionalNames'
+>
 
 export interface DropTypeNode extends OperationNode {
   readonly kind: 'DropTypeNode'
   readonly name: SchemableIdentifierNode
+  readonly additionalNames?: SchemableIdentifierNode[]
   readonly ifExists?: boolean
+  readonly cascade?: boolean
 }
 
 /**
@@ -18,10 +23,17 @@ export const DropTypeNode = freeze({
     return node.kind === 'DropTypeNode'
   },
 
-  create(name: SchemableIdentifierNode): DropTypeNode {
+  create(
+    names: SchemableIdentifierNode | SchemableIdentifierNode[],
+  ): DropTypeNode {
+    if (!Array.isArray(names)) {
+      names = [names]
+    }
+
     return freeze({
       kind: 'DropTypeNode',
-      name,
+      name: names[0],
+      additionalNames: names.slice(1),
     })
   },
 

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

@@ -1032,6 +1032,8 @@ export class OperationNodeTransformer {
     return requireAllProps<DropTypeNode>({
       kind: 'DropTypeNode',
       name: this.transformNode(node.name, queryId),
+      additionalNames: this.transformNodeList(node.additionalNames, queryId),
+      cascade: node.cascade,
       ifExists: node.ifExists,
     })
   }

src/parser/identifier-parser.ts

@@ -16,6 +16,16 @@ export function parseSchemableIdentifier(id: string): SchemableIdentifierNode {
   }
 }
 
+export function parseSchemableIdentifierArray(
+  id: string | string[],
+): SchemableIdentifierNode[] {
+  if (!Array.isArray(id)) {
+    id = [id]
+  }
+
+  return id.map(parseSchemableIdentifier)
+}
+
 function trim(str: string): string {
   return str.trim()
 }

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

@@ -1429,6 +1429,15 @@ export class DefaultQueryCompiler
     }
 
     this.visitNode(node.name)
+
+    if (node.additionalNames?.length) {
+      this.append(', ')
+      this.compileList(node.additionalNames)
+    }
+
+    if (node.cascade) {
+      this.append(' cascade')
+    }
   }
 
   protected override visitExplain(node: ExplainNode): void {

src/schema/drop-type-builder.ts

@@ -13,6 +13,9 @@ export class DropTypeBuilder implements OperationNodeSource, Compilable {
     this.#props = freeze(props)
   }
 
+  /**
+   * Adds `if exists` to the query.
+   */
   ifExists(): DropTypeBuilder {
     return new DropTypeBuilder({
       ...this.#props,
@@ -22,6 +25,18 @@ export class DropTypeBuilder implements OperationNodeSource, Compilable {
     })
   }
 
+  /**
+   * Adds `cascade` to the query.
+   */
+  cascade(): DropTypeBuilder {
+    return new DropTypeBuilder({
+      ...this.#props,
+      node: DropTypeNode.cloneWith(this.#props.node, {
+        cascade: true,
+      }),
+    })
+  }
+
   /**
    * Simply calls the provided function passing `this` as the only argument. `$call` returns
    * what the provided function returns.

src/schema/schema.ts

@@ -25,7 +25,10 @@ import { CreateTypeBuilder } from './create-type-builder.js'
 import { DropTypeBuilder } from './drop-type-builder.js'
 import { CreateTypeNode } from '../operation-node/create-type-node.js'
 import { DropTypeNode } from '../operation-node/drop-type-node.js'
-import { parseSchemableIdentifier } from '../parser/identifier-parser.js'
+import {
+  parseSchemableIdentifier,
+  parseSchemableIdentifierArray,
+} from '../parser/identifier-parser.js'
 import { RefreshMaterializedViewBuilder } from './refresh-materialized-view-builder.js'
 import { RefreshMaterializedViewNode } from '../operation-node/refresh-materialized-view-node.js'
 
@@ -311,12 +314,22 @@ export class SchemaModule {
    *   .ifExists()
    *   .execute()
    * ```
+   *
+   * You can also provide multiple type names:
+   *
+   * ```ts
+   * await db.schema
+   *   .dropType(['species', 'colors'])
+   *   .ifExists()
+   *   .cascade()
+   *   .execute()
+   * ```
    */
-  dropType(typeName: string): DropTypeBuilder {
+  dropType(typeName: string | string[]): DropTypeBuilder {
     return new DropTypeBuilder({
       queryId: createQueryId(),
       executor: this.#executor,
-      node: DropTypeNode.create(parseSchemableIdentifier(typeName)),
+      node: DropTypeNode.create(parseSchemableIdentifierArray(typeName)),
     })
   }
 

test/node/src/schema.test.ts

@@ -2457,6 +2457,7 @@ for (const dialect of DIALECTS) {
       if (sqlSpec === 'postgres') {
         it('should drop a schema cascade', async () => {
           await ctx.db.schema.createSchema('pets').execute()
+
           const builder = ctx.db.schema.dropSchema('pets').cascade()
 
           testSql(builder, dialect, {
@@ -2561,10 +2562,68 @@ for (const dialect of DIALECTS) {
 
           await builder.execute()
         })
+
+        it('should drop multiple types', async () => {
+          await ctx.db.schema.createType('species').execute()
+          await ctx.db.schema.createType('colors').execute()
+
+          const builder = ctx.db.schema
+            .dropType(['species', 'colors'])
+            .ifExists()
+
+          testSql(builder, dialect, {
+            postgres: {
+              sql: `drop type if exists "species", "colors"`,
+              parameters: [],
+            },
+            mysql: NOT_SUPPORTED,
+            mssql: NOT_SUPPORTED,
+            sqlite: NOT_SUPPORTED,
+          })
+
+          await builder.execute()
+        })
+
+        it('should drop multiple types if exists', async () => {
+          await ctx.db.schema.createType('species').execute()
+          const builder = ctx.db.schema
+            .dropType(['species', 'colors'])
+            .ifExists()
+
+          testSql(builder, dialect, {
+            postgres: {
+              sql: `drop type if exists "species", "colors"`,
+              parameters: [],
+            },
+            mysql: NOT_SUPPORTED,
+            mssql: NOT_SUPPORTED,
+            sqlite: NOT_SUPPORTED,
+          })
+
+          await builder.execute()
+        })
+
+        it('should drop a type and cascade', async () => {
+          await ctx.db.schema.createType('species').execute()
+
+          const builder = ctx.db.schema.dropType('species').cascade()
+
+          testSql(builder, dialect, {
+            postgres: {
+              sql: `drop type "species" cascade`,
+              parameters: [],
+            },
+            mysql: NOT_SUPPORTED,
+            mssql: NOT_SUPPORTED,
+            sqlite: NOT_SUPPORTED,
+          })
+
+          await builder.execute()
+        })
       }
 
       async function cleanup() {
-        await ctx.db.schema.dropType('species').ifExists().execute()
+        await ctx.db.schema.dropType(['species', 'colors']).ifExists().execute()
       }
     })