fix(typegen): schemas filtering

Author: avallete

src/server/templates/python.ts

@@ -17,23 +17,18 @@ export const apply = ({
   types,
 }: GeneratorMetadata): string => {
   const ctx = new PythonContext(types, columns, schemas)
-  const py_tables = tables
-    .filter((table) => schemas.some((schema) => schema.name === table.schema))
-    .flatMap((table) => {
-      const py_class_and_methods = ctx.tableToClass(table)
-      return py_class_and_methods
-    })
+  // Used for efficient lookup of types by schema name
+  const schemasNames = new Set(schemas.map((schema) => schema.name))
+  const py_tables = tables.flatMap((table) => {
+    const py_class_and_methods = ctx.tableToClass(table)
+    return py_class_and_methods
+  })
   const composite_types = types
-    .filter(
-      (type) => type.attributes.length > 0 && schemas.some((schema) => type.schema == schema.name)
-    )
+    // We always include system schemas, so we need to filter out types that are not in the included schemas
+    .filter((type) => type.attributes.length > 0 && schemasNames.has(type.schema))
     .map((type) => ctx.typeToClass(type))
-  const py_views = views
-    .filter((view) => schemas.some((schema) => schema.name === view.schema))
-    .map((view) => ctx.viewToClass(view))
-  const py_matviews = materializedViews
-    .filter((matview) => schemas.some((schema) => schema.name === matview.schema))
-    .map((matview) => ctx.matViewToClass(matview))
+  const py_views = views.map((view) => ctx.viewToClass(view))
+  const py_matviews = materializedViews.map((matview) => ctx.matViewToClass(matview))
 
   let output = `
 from __future__ import annotations

test/server/typegen.ts

@@ -6501,3 +6501,77 @@ class PublicCompositeTypeWithRecordAttribute(BaseModel):
     todo: PublicTodos = Field(alias="todo")"
 `)
 })
+
+test('typegen: python w/ excluded/included schemas', async () => {
+  // Create a test schema with some tables
+  await app.inject({
+    method: 'POST',
+    path: '/query',
+    payload: {
+      query: `
+        CREATE SCHEMA IF NOT EXISTS test_schema;
+        CREATE TABLE IF NOT EXISTS test_schema.test_table (
+          id serial PRIMARY KEY,
+          name text
+        );
+        CREATE TABLE IF NOT EXISTS test_schema.another_table (
+          id serial PRIMARY KEY,
+          value text
+        );
+      `,
+    },
+  })
+
+  try {
+    // Test excluded_schemas - should exclude test_schema
+    const { body: excludedBody } = await app.inject({
+      method: 'GET',
+      path: '/generators/python',
+      query: { access_control: 'public', excluded_schemas: 'test_schema' },
+    })
+    expect(excludedBody).not.toContain('TestSchemaTestTable')
+    expect(excludedBody).not.toContain('TestSchemaAnotherTable')
+    expect(excludedBody).toContain('PublicUsers')
+    expect(excludedBody).toContain('PublicTodos')
+
+    // Test included_schemas - should only include test_schema
+    const { body: includedBody } = await app.inject({
+      method: 'GET',
+      path: '/generators/python',
+      query: { access_control: 'public', included_schemas: 'test_schema' },
+    })
+    expect(includedBody).toContain('TestSchemaTestTable')
+    expect(includedBody).toContain('TestSchemaAnotherTable')
+    expect(includedBody).not.toContain('PublicUsers')
+    expect(includedBody).not.toContain('PublicTodos')
+
+    // Test multiple excluded schemas
+    const { body: multipleExcludedBody } = await app.inject({
+      method: 'GET',
+      path: '/generators/python',
+      query: { access_control: 'public', excluded_schemas: 'test_schema,public' },
+    })
+    expect(multipleExcludedBody).not.toContain('TestSchemaTestTable')
+    expect(multipleExcludedBody).not.toContain('PublicUsers')
+
+    // // Test multiple included schemas
+    const { body: multipleIncludedBody } = await app.inject({
+      method: 'GET',
+      path: '/generators/python',
+      query: { access_control: 'public', included_schemas: 'public,test_schema' },
+    })
+    expect(multipleIncludedBody).toContain('TestSchemaTestTable')
+    expect(multipleIncludedBody).toContain('PublicUsers')
+  } finally {
+    // Clean up test schema
+    await app.inject({
+      method: 'POST',
+      path: '/query',
+      payload: {
+        query: `
+          DROP SCHEMA IF EXISTS test_schema CASCADE;
+        `,
+      },
+    })
+  }
+})