aboutsummaryrefslogtreecommitdiff
path: root/src/hydrilla/proxy/state_impl/base.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/hydrilla/proxy/state_impl/base.py')
-rw-r--r--src/hydrilla/proxy/state_impl/base.py24
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)