Support tables in different schemas

L-Mario564 · Aug 22, 2024 · f5c2145 · f5c2145
1 parent 32e18bc
commit f5c2145
Show file tree

Hide file tree

Showing 7 changed files with 187 additions and 13 deletions.
diff --git a/src/__tests__/mysql/_mysql.test.ts b/src/__tests__/mysql/_mysql.test.ts
@@ -16,6 +16,7 @@ import {
   int,
   mediumint,
   mysqlEnum,
+  mysqlSchema,
   mysqlTable,
   primaryKey,
   real,
@@ -165,6 +166,25 @@ async function indexesTest() {
   expect(result).toBe(true);
 }
 
+async function schemasTest() {
+  const schema1 = mysqlSchema('schema1');
+  const schema2 = mysqlSchema('schema2');
+  const table1 = schema1.table('table1', {
+    id: serial('id').primaryKey(),
+    tabl2Id: int('table2_id').references(() => table2.id)
+  });
+  const table2 = schema2.table('table2', {
+    id: serial('id').primaryKey()
+  });
+
+  const schema = { schema1, schema2, table1, table2 };
+  const out = `${pathPrefix}schemas.generated.dbml`;
+  mysqlGenerate({ schema, out });
+
+  const result = await compareContents(out);
+  expect(result).toBe(true);
+}
+
 async function rqbTest() {
   const users = mysqlTable('users', {
     id: serial('id').primaryKey(),
@@ -218,6 +238,34 @@ async function rqbTest() {
   expect(result).toBe(true);
 }
 
+async function schemasRQBTest() {
+  const schema1 = mysqlSchema('schema1');
+  const schema2 = mysqlSchema('schema2');
+  const table1 = schema1.table('table1', {
+    id: serial('id').primaryKey(),
+    tabl2Id: int('table2_id').references(() => table2.id)
+  });
+  const table1Relations = relations(table1, ({ one }) => ({
+    table2: one(table2, {
+      fields: [table1.tabl2Id],
+      references: [table2.id]
+    })
+  }));
+  const table2 = schema2.table('table2', {
+    id: serial('id').primaryKey()
+  });
+  const table2Relations = relations(table2, ({ many }) => ({
+    table1: many(table1)
+  }));
+
+  const schema = { schema1, schema2, table1, table1Relations, table2, table2Relations };
+  const out = `${pathPrefix}schemas-rqb.generated.dbml`;
+  mysqlGenerate({ schema, out });
+
+  const result = await compareContents(out);
+  expect(result).toBe(true);
+}
+
 async function realTest() {
   const users = mysqlTable('users', {
     id: serial('id').primaryKey(),
@@ -327,6 +375,8 @@ describe('MySQL dialect tests', () => {
   it('Outputs an inline foreign key', inlineFkTest);
   it('Outputs a foreign key', fkTest);
   it('Outputs all indexes', indexesTest);
+  it('Outputs tables with different schemas', schemasTest);
   it('Outputs relations written with the RQB API', rqbTest);
+  it('Outputs tables with different schemas with the RQB API', schemasRQBTest);
   it('Outputs the result of a more "realistic" schema', realTest);
 });
diff --git a/src/__tests__/mysql/schemas-rqb.dbml b/src/__tests__/mysql/schemas-rqb.dbml
@@ -0,0 +1,10 @@
+table schema1.table1 {
+  id serial [pk, not null, increment]
+  table2_id int
+}
+
+table schema2.table2 {
+  id serial [pk, not null, increment]
+}
+
+ref table1_table2_id_table2_id_fk: schema1.table1.table2_id > schema2.table2.id [delete: no action, update: no action]
diff --git a/src/__tests__/mysql/schemas.dbml b/src/__tests__/mysql/schemas.dbml
@@ -0,0 +1,10 @@
+table schema1.table1 {
+  id serial [pk, not null, increment]
+  table2_id int
+}
+
+table schema2.table2 {
+  id serial [pk, not null, increment]
+}
+
+ref table1_table2_id_table2_id_fk: schema1.table1.table2_id > schema2.table2.id [delete: no action, update: no action]
diff --git a/src/__tests__/pg/_pg.test.ts b/src/__tests__/pg/_pg.test.ts
@@ -14,6 +14,7 @@ import {
   jsonb,
   numeric,
   pgEnum,
+  pgSchema,
   pgTable,
   primaryKey,
   real,
@@ -175,6 +176,25 @@ async function indexesTest() {
   expect(result).toBe(true);
 }
 
+async function schemasTest() {
+  const schema1 = pgSchema('schema1');
+  const schema2 = pgSchema('schema2');
+  const table1 = schema1.table('table1', {
+    id: serial('id').primaryKey(),
+    tabl2Id: integer('table2_id').references(() => table2.id)
+  });
+  const table2 = schema2.table('table2', {
+    id: serial('id').primaryKey()
+  });
+
+  const schema = { schema1, schema2, table1, table2 };
+  const out = `${pathPrefix}schemas.generated.dbml`;
+  pgGenerate({ schema, out });
+
+  const result = await compareContents(out);
+  expect(result).toBe(true);
+}
+
 async function rqbTest() {
   const users = pgTable('users', {
     id: serial('id').primaryKey(),
@@ -228,6 +248,34 @@ async function rqbTest() {
   expect(result).toBe(true);
 }
 
+async function schemasRQBTest() {
+  const schema1 = pgSchema('schema1');
+  const schema2 = pgSchema('schema2');
+  const table1 = schema1.table('table1', {
+    id: serial('id').primaryKey(),
+    tabl2Id: integer('table2_id').references(() => table2.id)
+  });
+  const table1Relations = relations(table1, ({ one }) => ({
+    table2: one(table2, {
+      fields: [table1.tabl2Id],
+      references: [table2.id]
+    })
+  }));
+  const table2 = schema2.table('table2', {
+    id: serial('id').primaryKey()
+  });
+  const table2Relations = relations(table2, ({ many }) => ({
+    table1: many(table1)
+  }));
+
+  const schema = { schema1, schema2, table1, table1Relations, table2, table2Relations };
+  const out = `${pathPrefix}schemas-rqb.generated.dbml`;
+  pgGenerate({ schema, out });
+
+  const result = await compareContents(out);
+  expect(result).toBe(true);
+}
+
 async function realTest() {
   const users = pgTable('users', {
     id: serial('id').primaryKey(),
@@ -337,6 +385,8 @@ describe('Postgres dialect tests', () => {
   it('Outputs an inline foreign key', inlineFkTest);
   it('Outputs a foreign key', fkTest);
   it('Outputs all indexes', indexesTest);
+  it('Outputs tables with different schemas', schemasTest);
   it('Outputs relations written with the RQB API', rqbTest);
+  it('Outputs tables with different schemas with the RQB API', schemasRQBTest);
   it('Outputs the result of a more "realistic" schema', realTest);
 });
diff --git a/src/__tests__/pg/schemas-rqb.dbml b/src/__tests__/pg/schemas-rqb.dbml
@@ -0,0 +1,10 @@
+table schema1.table1 {
+  id serial [pk, not null, increment]
+  table2_id integer
+}
+
+table schema2.table2 {
+  id serial [pk, not null, increment]
+}
+
+ref table1_table2_id_table2_id_fk: schema1.table1.table2_id > schema2.table2.id [delete: no action, update: no action]
diff --git a/src/__tests__/pg/schemas.dbml b/src/__tests__/pg/schemas.dbml
@@ -0,0 +1,10 @@
+table schema1.table1 {
+  id serial [pk, not null, increment]
+  table2_id integer
+}
+
+table schema2.table2 {
+  id serial [pk, not null, increment]
+}
+
+ref table1_table2_id_table2_id_fk: schema1.table1.table2_id > schema2.table2.id [delete: no action, update: no action]
diff --git a/src/generators/common.ts b/src/generators/common.ts
@@ -221,15 +221,33 @@ export abstract class BaseGenerator<
 
   private generateForeignKeys(fks: (PgForeignKey | MySqlForeignKey | SQLiteForeignKey)[]) {
     for (let i = 0; i < fks.length; i++) {
-      const dbml = new DBML()
-        .insert(`ref ${fks[i].getName()}: `)
-        .escapeSpaces((fks[i].table as unknown as AnyTable)[TableName])
+      const sourceTable = fks[i].table as unknown as AnyTable;
+      const foreignTable = fks[i].reference().foreignTable as unknown as AnyTable;
+      const sourceSchema = sourceTable[Schema];
+      const foreignSchema = foreignTable[Schema];
+      const sourceColumns = fks[i].reference().columns;
+      const foreignColumns = fks[i].reference().foreignColumns;
+
+      const dbml = new DBML().insert(`ref ${fks[i].getName()}: `);
+
+      if (sourceSchema) {
+        dbml.escapeSpaces(sourceSchema).insert('.');
+      }
+
+      dbml
+        .escapeSpaces(sourceTable[TableName])
         .insert('.')
-        .insert(wrapColumns(fks[i].reference().columns, this.buildQueryConfig.escapeName))
-        .insert(' > ')
-        .escapeSpaces((fks[i].reference().foreignTable as unknown as AnyTable)[TableName])
+        .insert(wrapColumns(sourceColumns, this.buildQueryConfig.escapeName))
+        .insert(' > ');
+
+      if (foreignSchema) {
+        dbml.escapeSpaces(foreignSchema).insert('.');
+      }
+
+      dbml
+        .escapeSpaces(foreignTable[TableName])
         .insert('.')
-        .insert(wrapColumns(fks[i].reference().foreignColumns, this.buildQueryConfig.escapeName));
+        .insert(wrapColumns(foreignColumns, this.buildQueryConfig.escapeName));
 
       const actions: string[] = [
         `delete: ${fks[i].onDelete || 'no action'}`,
@@ -247,8 +265,10 @@ export abstract class BaseGenerator<
       string,
       {
         type: 'one' | 'many';
+        sourceSchema?: string;
         sourceTable?: string;
         sourceColumns?: string[];
+        foreignSchema?: string;
         foreignTable?: string;
         foreignColumns?: string[];
       }
@@ -274,8 +294,10 @@ export abstract class BaseGenerator<
         if ((is(relation, One) && relation.config?.references.length) || 0 > 0) {
           left[key] = {
             type: 'one',
+            sourceSchema: (relation.sourceTable as unknown as AnyTable)[Schema],
             sourceTable: (relation.sourceTable as unknown as AnyTable)[TableName],
             sourceColumns: (relation as One).config?.fields.map((col) => col.name) || [],
+            foreignSchema: (relation.referencedTable as unknown as AnyTable)[Schema],
             foreignTable: relation.referencedTableName,
             foreignColumns: (relation as One).config?.references.map((col) => col.name) || []
           };
@@ -288,7 +310,9 @@ export abstract class BaseGenerator<
     }
 
     for (const key in left) {
+      const sourceSchema = left[key].sourceSchema || '';
       const sourceTable = left[key].sourceTable || '';
+      const foreignSchema = left[key].foreignSchema || '';
       const foreignTable = left[key].foreignTable || '';
       const sourceColumns = left[key].sourceColumns || [];
       const foreignColumns = left[key].foreignColumns || [];
@@ -300,18 +324,28 @@ export abstract class BaseGenerator<
         );
       }
 
-      const dbml = new DBML()
-        .insert('ref: ')
+      const dbml = new DBML().insert('ref: ');
+
+      if (sourceSchema) {
+        dbml.escapeSpaces(sourceSchema).insert('.');
+      }
+
+      dbml
         .escapeSpaces(sourceTable)
         .insert('.')
         .insert(wrapColumnNames(sourceColumns, this.buildQueryConfig.escapeName))
-        .insert(` ${relationType === 'one' ? '-' : '>'} `)
+        .insert(` ${relationType === 'one' ? '-' : '>'} `);
+
+      if (foreignSchema) {
+        dbml.escapeSpaces(foreignSchema).insert('.');
+      }
+
+      dbml
         .escapeSpaces(foreignTable)
         .insert('.')
-        .insert(wrapColumnNames(foreignColumns, this.buildQueryConfig.escapeName))
-        .build();
+        .insert(wrapColumnNames(foreignColumns, this.buildQueryConfig.escapeName));
 
-      this.generatedRefs.push(dbml);
+      this.generatedRefs.push(dbml.build());
     }
   }