fix(schema): support complex Drizzle method chaining patterns

.changeset/whole-mirrors-throw.md

@@ -0,0 +1,5 @@
+---
+"@liam-hq/schema": patch
+---
+
+🐛 fix(Drizzle parser): support chained methods like pgTable().enableRLS().$comment()

frontend/packages/schema/src/parser/drizzle/postgres/__tests__/index.test.ts

@@ -57,5 +57,64 @@ describe(processor, () => {
         'numeric(10,2)',
       )
     })
+
+    it('should parse tables with .enableRLS() method chaining', async () => {
+      const { value } = await processor(`
+        import { pgTable, serial, text, varchar } from 'drizzle-orm/pg-core';
+
+        export const users = pgTable('users', {
+          id: serial('id').primaryKey(),
+          name: varchar('name', { length: 255 }).notNull(),
+          email: text('email').notNull(),
+        }).enableRLS();
+      `)
+
+      // Verify that table with .enableRLS() is correctly parsed
+      expect(Object.keys(value.tables)).toHaveLength(1)
+      expect(value.tables['users']).toBeDefined()
+
+      // Verify table structure
+      expect(value.tables['users']?.columns['id']?.type).toBe('serial')
+      expect(value.tables['users']?.columns['name']?.type).toBe('varchar(255)')
+      expect(value.tables['users']?.columns['email']?.type).toBe('text')
+    })
+
+    it('should parse tables with .$comment() method chaining', async () => {
+      const { value } = await processor(`
+        import { pgTable, serial, text, varchar } from 'drizzle-orm/pg-core';
+
+        export const users = pgTable('users', {
+          id: serial('id').primaryKey(),
+          name: varchar('name', { length: 255 }).notNull(),
+          email: text('email').notNull(),
+        }).$comment('User accounts table');
+      `)
+
+      // Verify that table with .$comment() is correctly parsed
+      expect(Object.keys(value.tables)).toHaveLength(1)
+      expect(value.tables['users']).toBeDefined()
+
+      // Verify table structure and comment
+      expect(value.tables['users']?.comment).toBe('User accounts table')
+      expect(value.tables['users']?.columns['name']?.type).toBe('varchar(255)')
+      expect(value.tables['users']?.columns['email']?.type).toBe('text')
+    })
+
+    it('should parse tables with complex method chaining (.enableRLS() + .$comment())', async () => {
+      const { value } = await processor(`
+        import { pgTable, serial, text } from 'drizzle-orm/pg-core';
+
+        export const mixed = pgTable('mixed', {
+          id: serial('id').primaryKey(),
+          data: text('data'),
+        }).enableRLS().$comment('Complex chaining example');
+      `)
+
+      // Verify that table with complex chaining is correctly parsed
+      expect(Object.keys(value.tables)).toHaveLength(1)
+      expect(value.tables['mixed']).toBeDefined()
+      expect(value.tables['mixed']?.comment).toBe('Complex chaining example')
+      expect(value.tables['mixed']?.columns['data']?.type).toBe('text')
+    })
   })
 })

frontend/packages/schema/src/parser/drizzle/postgres/astUtils.ts

@@ -111,9 +111,32 @@ export const isMemberExpression = (
 /**
  * Check if a call expression is a pgTable call
  */
-export const isPgTableCall = (callExpr: CallExpression): boolean => {
+const isPgTableCall = (callExpr: CallExpression): boolean => {
   return isIdentifierWithName(callExpr.callee, 'pgTable')
 }
+/**
+ * Extract the base pgTable call from method chaining patterns
+ * Handles patterns like: pgTable(...).enableRLS(), pgTable(...).comment(...), etc.
+ */
+export const extractPgTableFromChain = (
+  callExpr: CallExpression,
+): CallExpression | null => {
+  // If it's already a direct pgTable call, return it
+  if (isPgTableCall(callExpr)) {
+    return callExpr
+  }
+
+  // If it's a method call on another expression, check the object
+  if (callExpr.callee.type === 'MemberExpression') {
+    const baseCall = callExpr.callee.object
+    if (baseCall.type === 'CallExpression') {
+      // Recursively check if the base call is or contains a pgTable call
+      return extractPgTableFromChain(baseCall)
+    }
+  }
+
+  return null
+}
 
 /**
  * Check if a call expression is a schema.table() call

frontend/packages/schema/src/parser/drizzle/postgres/mainParser.ts

@@ -2,10 +2,10 @@
  * Main orchestrator for Drizzle ORM schema parsing
  */
 
-import type { Module, VariableDeclarator } from '@swc/core'
+import type { CallExpression, Module, VariableDeclarator } from '@swc/core'
 import { parseSync } from '@swc/core'
 import type { Processor, ProcessResult } from '../../types.js'
-import { isPgTableCall, isSchemaTableCall } from './astUtils.js'
+import { extractPgTableFromChain, isSchemaTableCall } from './astUtils.js'
 import {
   convertDrizzleEnumsToInternal,
   convertDrizzleTablesToInternal,
@@ -80,53 +80,118 @@ const visitModule = (
 }
 
 /**
- * Visit variable declarator to find pgTable, pgEnum, or relations calls
+ * Check if the call expression is a comment call
  */
-const visitVariableDeclarator = (
+const isCommentCall = (callExpr: CallExpression): boolean => {
+  return (
+    callExpr.type === 'CallExpression' &&
+    callExpr.callee.type === 'MemberExpression' &&
+    callExpr.callee.property.type === 'Identifier' &&
+    callExpr.callee.property.value === '$comment'
+  )
+}
+
+/**
+ * Check if the call expression is a pgEnum call
+ */
+const isPgEnumCall = (callExpr: CallExpression): boolean => {
+  return (
+    callExpr.callee.type === 'Identifier' && callExpr.callee.value === 'pgEnum'
+  )
+}
+
+/**
+ * Handle comment calls
+ */
+const handleCommentCall = (
   declarator: VariableDeclarator,
   tables: Record<string, DrizzleTableDefinition>,
-  enums: Record<string, DrizzleEnumDefinition>,
   variableToTableMapping: Record<string, string>,
 ) => {
-  if (!declarator.init || declarator.init.type !== 'CallExpression') return
+  if (declarator.init?.type !== 'CallExpression') return
 
-  const callExpr = declarator.init
+  const table = parsePgTableWithComment(declarator.init)
+  if (table && declarator.id.type === 'Identifier') {
+    tables[table.name] = table
+    variableToTableMapping[declarator.id.value] = table.name
+  }
+}
+
+/**
+ * Handle schema table calls
+ */
+const handleSchemaTableCall = (
+  declarator: VariableDeclarator,
+  callExpr: CallExpression,
+  tables: Record<string, DrizzleTableDefinition>,
+  variableToTableMapping: Record<string, string>,
+) => {
+  const table = parseSchemaTableCall(callExpr)
+  if (table && declarator.id.type === 'Identifier') {
+    tables[table.name] = table
+    variableToTableMapping[declarator.id.value] = table.name
+  }
+}
+
+/**
+ * Handle pgEnum calls
+ */
+const handlePgEnumCall = (
+  declarator: VariableDeclarator,
+  callExpr: CallExpression,
+  enums: Record<string, DrizzleEnumDefinition>,
+) => {
+  const enumDef = parsePgEnumCall(callExpr)
+  if (enumDef && declarator.id.type === 'Identifier') {
+    const variableName = declarator.id.value
+    // Only store by variable name to avoid conflicts between actual name and variable name
+    enums[variableName] = enumDef
+  }
+}
 
-  if (isPgTableCall(callExpr)) {
-    const table = parsePgTableCall(callExpr)
+/**
+ * Handle pgTable calls (direct or method chained)
+ */
+const handlePgTableCall = (
+  declarator: VariableDeclarator,
+  callExpr: CallExpression,
+  tables: Record<string, DrizzleTableDefinition>,
+  variableToTableMapping: Record<string, string>,
+) => {
+  const basePgTableCall = extractPgTableFromChain(callExpr)
+  if (basePgTableCall) {
+    const table = parsePgTableCall(basePgTableCall)
     if (table && declarator.id.type === 'Identifier') {
       tables[table.name] = table
-      // Map variable name to table name
       variableToTableMapping[declarator.id.value] = table.name
     }
+  }
+}
+
+/**
+ * Visit variable declarator to find pgTable, pgEnum, or relations calls
+ */
+const visitVariableDeclarator = (
+  declarator: VariableDeclarator,
+  tables: Record<string, DrizzleTableDefinition>,
+  enums: Record<string, DrizzleEnumDefinition>,
+  variableToTableMapping: Record<string, string>,
+) => {
+  if (!declarator.init || declarator.init.type !== 'CallExpression') {
+    return
+  }
+
+  const callExpr = declarator.init
+
+  // Handle different types of call expressions
+  if (isCommentCall(declarator.init)) {
+    handleCommentCall(declarator, tables, variableToTableMapping)
   } else if (isSchemaTableCall(callExpr)) {
-    const table = parseSchemaTableCall(callExpr)
-    if (table && declarator.id.type === 'Identifier') {
-      tables[table.name] = table
-      // Map variable name to table name
-      variableToTableMapping[declarator.id.value] = table.name
-    }
-  } else if (
-    declarator.init.type === 'CallExpression' &&
-    declarator.init.callee.type === 'MemberExpression' &&
-    declarator.init.callee.property.type === 'Identifier' &&
-    declarator.init.callee.property.value === '$comment'
-  ) {
-    // Handle table comments: pgTable(...).comment(...)
-    const table = parsePgTableWithComment(declarator.init)
-    if (table && declarator.id.type === 'Identifier') {
-      tables[table.name] = table
-      // Map variable name to table name
-      variableToTableMapping[declarator.id.value] = table.name
-    }
-  } else if (
-    callExpr.callee.type === 'Identifier' &&
-    callExpr.callee.value === 'pgEnum'
-  ) {
-    const enumDef = parsePgEnumCall(callExpr)
-    if (enumDef && declarator.id.type === 'Identifier') {
-      enums[declarator.id.value] = enumDef
-    }
+    handleSchemaTableCall(declarator, callExpr, tables, variableToTableMapping)
+  } else if (isPgEnumCall(callExpr)) {
+    handlePgEnumCall(declarator, callExpr, enums)
+  } else {
+    handlePgTableCall(declarator, callExpr, tables, variableToTableMapping)
   }
 }
 

frontend/packages/schema/src/parser/drizzle/postgres/tableParser.ts

@@ -4,11 +4,11 @@
 
 import type { CallExpression, Expression } from '@swc/core'
 import {
+  extractPgTableFromChain,
   getArgumentExpression,
   getIdentifierName,
   getStringValue,
   isObjectExpression,
-  isPgTableCall,
   isSchemaTableCall,
   isStringLiteral,
 } from './astUtils.js'
@@ -37,15 +37,19 @@ export const parsePgTableWithComment = (
     }
   }
 
-  // Get the pgTable call from the object of the member expression
+  // Get the pgTable call from the chain - handle complex chaining like pgTable().enableRLS().$comment()
   if (commentCallExpr.callee.type === 'MemberExpression') {
-    const pgTableCall = commentCallExpr.callee.object
-    if (pgTableCall.type === 'CallExpression' && isPgTableCall(pgTableCall)) {
-      const table = parsePgTableCall(pgTableCall)
-      if (table && comment) {
-        table.comment = comment
+    const chainCall = commentCallExpr.callee.object
+    if (chainCall.type === 'CallExpression') {
+      // Use extractPgTableFromChain to handle complex method chaining
+      const pgTableCall = extractPgTableFromChain(chainCall)
+      if (pgTableCall) {
+        const table = parsePgTableCall(pgTableCall)
+        if (table && comment) {
+          table.comment = comment
+        }
+        return table
       }
-      return table
     }
   }