diff options
Diffstat (limited to 'src/hydrilla/proxy/state_impl/base.py')
-rw-r--r-- | src/hydrilla/proxy/state_impl/base.py | 24 |
1 files changed, 14 insertions, 10 deletions
diff --git a/src/hydrilla/proxy/state_impl/base.py b/src/hydrilla/proxy/state_impl/base.py index e5a9898..82b8734 100644 --- a/src/hydrilla/proxy/state_impl/base.py +++ b/src/hydrilla/proxy/state_impl/base.py @@ -50,22 +50,26 @@ from .. import policies @contextmanager -def temporary_ids_table( - cursor: sqlite3.Cursor, - ids: t.Iterable[int], - table_name: str = '__helper_ids' +def temporary_ids_tables( + cursor: sqlite3.Cursor, + tables: t.Iterable[tuple[str, t.Iterable[int]]] ) -> t.Iterator[None]: - cursor.execute( - f'CREATE TEMPORARY TABLE "{table_name}"(id INTEGER PRIMARY KEY);' - ) + created: set[str] = set() try: - for id in ids: - cursor.execute(f'INSERT INTO "{table_name}" VALUES(?);', (id,)) + for name, ids in tables: + cursor.execute( + f'CREATE TEMPORARY TABLE "{name}"(id INTEGER PRIMARY KEY);' + ) + created.add(name) + + for id in ids: + cursor.execute(f'INSERT INTO "{name}" VALUES(?);', (id,)) yield finally: - cursor.execute(f'DROP TABLE "{table_name}";') + for name in created: + cursor.execute(f'DROP TABLE "{name}";') @dc.dataclass(frozen=True) |