diff --git a/src/server/templates/python.ts b/src/server/templates/python.ts index 100870b4..fa2d5eef 100644 --- a/src/server/templates/python.ts +++ b/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 diff --git a/test/server/typegen.ts b/test/server/typegen.ts index 46079de0..2005339e 100644 --- a/test/server/typegen.ts +++ b/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; + `, + }, + }) + } +})