Skip to content

Commit

Permalink
be able to pass an engine into TableStorage (#1109)
Browse files Browse the repository at this point in the history
  • Loading branch information
dantownsend authored Oct 18, 2024
1 parent 67be2d6 commit e2fa8f8
Showing 1 changed file with 17 additions and 5 deletions.
22 changes: 17 additions & 5 deletions piccolo/table_reflection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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.
"""
Expand Down

0 comments on commit e2fa8f8

Please sign in to comment.