diff --git a/piccolo/table_reflection.py b/piccolo/table_reflection.py index 69c0bb7e..3d050099 100644 --- a/piccolo/table_reflection.py +++ b/piccolo/table_reflection.py @@ -8,6 +8,8 @@ from dataclasses import dataclass from piccolo.apps.schema.commands.generate import get_output_schema +from piccolo.engine import engine_finder +from piccolo.engine.base import Engine from piccolo.table import Table @@ -78,9 +80,16 @@ class TableStorage(metaclass=Singleton): works with Postgres. """ - def __init__(self): + def __init__(self, engine: t.Optional[Engine] = None): + """ + :param engine: + Which engine to use to make the database queries. If not specified, + we try importing an engine from ``piccolo_conf.py``. + + """ + self.engine = engine or engine_finder() self.tables = ImmutableDict() - self._schema_tables = {} + self._schema_tables: t.Dict[str, t.List[str]] = {} async def reflect( self, @@ -120,10 +129,13 @@ async def reflect( exclude_list = self._to_list(exclude) if keep_existing: - exclude += self._schema_tables.get(schema_name, []) + exclude_list += self._schema_tables.get(schema_name, []) output_schema = await get_output_schema( - schema_name=schema_name, include=include_list, exclude=exclude_list + schema_name=schema_name, + include=include_list, + exclude=exclude_list, + engine=self.engine, ) add_tables = [ self._add_table(schema_name=schema_name, table=table) @@ -177,7 +189,7 @@ async def _add_table(self, schema_name: str, table: t.Type[Table]) -> None: def _add_to_schema_tables(self, schema_name: str, table_name: str) -> None: """ - We keep record of schemas and their tables for easy use. This method + We keep a record of schemas and their tables for easy use. This method adds a table to its schema. """