add raw and index type

faasjs · Feb 8, 2025 · b42d173 · b42d173
1 parent 67cfe2b
commit b42d173
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 10 deletions.
diff --git a/src/schema-builder/__tests__/index.test.ts b/src/schema-builder/__tests__/index.test.ts
@@ -125,6 +125,11 @@ describe('SchemaBuilder', () => {
         unique: true,
         check: 'updated_at > now()',
       })
+      table.index('number', {
+        unique: true,
+        indexType: 'btree',
+      })
+      table.raw('SELECT 1')
     })
 
     await builder.run()
@@ -158,5 +163,16 @@ describe('SchemaBuilder', () => {
       tablespace: null,
       indexdef: 'CREATE UNIQUE INDEX alters_pkey ON public.alters USING btree (updated_at)',
     })
+
+    expect(
+      indices.find(i => i.indexname.includes('idx_alters_number'))
+    ).toMatchObject({
+      schemaname: 'public',
+      tablename: 'alters',
+      indexname: 'idx_alters_number',
+      tablespace: null,
+      indexdef:
+        'CREATE UNIQUE INDEX idx_alters_number ON public.alters USING btree (number)',
+    })
   })
 })
diff --git a/src/schema-builder/table-builder.ts b/src/schema-builder/table-builder.ts
@@ -78,12 +78,19 @@ type AlterOperation =
 
 export type TableBuilderMode = 'create' | 'alter'
 
+export type IndexDefs = {
+  columns: string[]
+  unique?: boolean
+  indexType?: string
+}
+
 export class TableBuilder {
   private tableName: string
   private mode: TableBuilderMode
   private columns: Map<string, ColumnDefinition> = new Map()
-  private indices: Map<string, string[]> = new Map()
+  private indices: Map<string, IndexDefs> = new Map()
   private operations: AlterOperation[] = [];
+  private raws: string[] = []
 
   constructor(tableName: string, mode: TableBuilderMode) {
     this.tableName = tableName
@@ -199,16 +206,19 @@ export class TableBuilder {
     return this
   }
 
-  index(columns: string | string[], indexName?: string) {
-    const name = indexName || `idx_${this.tableName}_${Array.isArray(columns) ? columns.join('_') : columns}`
+  index(columns: string | string[], options: Omit<IndexDefs, 'columns'> = {}) {
+    const name = `idx_${this.tableName}_${Array.isArray(columns) ? columns.join('_') : columns}`
 
-    this.indices.set(name, Array.isArray(columns) ? columns : [columns])
+    this.indices.set(name, {
+      columns: Array.isArray(columns) ? columns : [columns],
+      ...options
+    })
 
     return this
   }
 
-  dropIndex(columns: string | string[], indexName?: string) {
-    const name = indexName || `idx_${this.tableName}_${Array.isArray(columns) ? columns.join('_') : columns}`
+  dropIndex(columns: string | string[]) {
+    const name = `idx_${this.tableName}_${Array.isArray(columns) ? columns.join('_') : columns}`
 
     if (this.indices.has(name)) {
       this.indices.delete(name)
@@ -224,6 +234,12 @@ export class TableBuilder {
     });
   }
 
+  raw(sql: string) {
+    this.raws.push(sql)
+
+    return this
+  }
+
   toSQL(): string[] {
     const sql: string[] = []
     const columnDefs = Array.from(this.columns.entries())
@@ -267,11 +283,11 @@ export class TableBuilder {
     }
 
     const indexDefs = Array.from(this.indices.entries())
-      .map(([name, columns]) => this.indexToSQL(name, columns))
+      .map(([name, definition]) => this.indexToSQL(name, definition))
 
     sql.push(indexDefs.join('\n'))
 
-    // console.log(sql.join('\n'))
+    sql.push(...this.raws)
 
     return sql
   }
@@ -291,8 +307,19 @@ export class TableBuilder {
     return parts.join(' ')
   }
 
-  private indexToSQL(name: string, columns: string[]): string {
-    return `CREATE INDEX ${escapeIdentifier(name)} ON ${escapeIdentifier(this.tableName)} (${columns.map(escapeIdentifier).join(', ')});`
+  private indexToSQL(name: string, defs: IndexDefs): string {
+    const parts = [
+      'CREATE',
+      defs.unique ? 'UNIQUE' : '',
+      'INDEX',
+      escapeIdentifier(name),
+      'ON',
+      escapeIdentifier(this.tableName),
+      defs.indexType ? `USING ${defs.indexType}` : '',
+      `(${defs.columns.map(escapeIdentifier).join(', ')});`
+    ].filter(Boolean)
+
+    return parts.join(' ')
   }
 
   private alterToSql(columnName: string, type: keyof ColumnDefinition, value: any) {