From 829766a624539714d2e11c40569059de2719570b Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Wed, 22 Jan 2025 09:16:33 +0200 Subject: [PATCH 1/4] Added pragma on file connections and DROP TABLE IF EXISTS --- ...2025_01_21_0820-4dec3e456c9e_add_on_delete_cascade.py | 6 +++++- src/codegate/db/connection.py | 9 ++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/migrations/versions/2025_01_21_0820-4dec3e456c9e_add_on_delete_cascade.py b/migrations/versions/2025_01_21_0820-4dec3e456c9e_add_on_delete_cascade.py index 5aaa4467..99a98404 100644 --- a/migrations/versions/2025_01_21_0820-4dec3e456c9e_add_on_delete_cascade.py +++ b/migrations/versions/2025_01_21_0820-4dec3e456c9e_add_on_delete_cascade.py @@ -21,6 +21,7 @@ def upgrade() -> None: # To add ON DELETE CASCADE to the foreign key constraint, we need to # rename the table, create a new table with the constraint, and copy # the data over. + op.execute("DROP TABLE IF EXISTS _prompts_old;") op.execute("ALTER TABLE prompts RENAME TO _prompts_old;") op.execute( """ @@ -39,6 +40,7 @@ def upgrade() -> None: op.execute("DROP TABLE _prompts_old;") # Doing the same for the sessions table + op.execute("DROP TABLE IF EXISTS _sessions_old;") op.execute("ALTER TABLE sessions RENAME TO _sessions_old;") op.execute( """ @@ -54,6 +56,7 @@ def upgrade() -> None: op.execute("DROP TABLE _sessions_old;") # Doing the same for the output table + op.execute("DROP TABLE IF EXISTS _outputs_old;") op.execute("ALTER TABLE outputs RENAME TO _outputs_old;") op.execute( """ @@ -70,6 +73,7 @@ def upgrade() -> None: op.execute("DROP TABLE _outputs_old;") # Doing the same for the alerts table + op.execute("DROP TABLE IF EXISTS _alerts_old;") op.execute("ALTER TABLE alerts RENAME TO _alerts_old;") op.execute( """ @@ -89,7 +93,7 @@ def upgrade() -> None: op.execute("DROP TABLE _alerts_old;") # Dropping unused table - op.execute("DROP TABLE settings;") + op.execute("DROP TABLE IF EXISTS settings;") # Create indexes for foreign keys op.execute("CREATE INDEX idx_outputs_prompt_id ON outputs(prompt_id);") diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index dc75bb83..91bd58b1 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -45,9 +45,11 @@ def set_sqlite_pragma(dbapi_connection, connection_record): [SQLite docs](https://www.sqlite.org/foreignkeys.html) [SO](https://stackoverflow.com/questions/2614984/sqlite-sqlalchemy-how-to-enforce-foreign-keys) """ - cursor = dbapi_connection.cursor() - cursor.execute("PRAGMA foreign_keys=ON") - cursor.close() + # Check if the connection is the one you want to apply the PRAGMA to + if connection_record.info.get('enable_fks', False): + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() class DbCodeGate: @@ -78,6 +80,7 @@ def __init__(self, sqlite_path: Optional[str] = None): "isolation_level": "AUTOCOMMIT", # Required for SQLite } self._async_db_engine = create_async_engine(**engine_dict) + self._async_db_engine.connect().connection.connection_record.info['enable_fks'] = True def does_db_exist(self): return self._db_path.is_file() From b4ce9d3578f2cf2fb8b7421ebef3b6d5794abfb1 Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Wed, 22 Jan 2025 09:18:03 +0200 Subject: [PATCH 2/4] Changed comment --- src/codegate/db/connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index 91bd58b1..8646ddc8 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -45,7 +45,7 @@ def set_sqlite_pragma(dbapi_connection, connection_record): [SQLite docs](https://www.sqlite.org/foreignkeys.html) [SO](https://stackoverflow.com/questions/2614984/sqlite-sqlalchemy-how-to-enforce-foreign-keys) """ - # Check if the connection is the one you want to apply the PRAGMA to + # Only enable foreign keys if the connection record has the flag set to True if connection_record.info.get('enable_fks', False): cursor = dbapi_connection.cursor() cursor.execute("PRAGMA foreign_keys=ON") From ef200a31773ada26bc8feeaa57d8f6fa6bcf9e5a Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Wed, 22 Jan 2025 09:19:32 +0200 Subject: [PATCH 3/4] Formatting changes --- src/codegate/db/connection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index 8646ddc8..5f290ffe 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -46,7 +46,7 @@ def set_sqlite_pragma(dbapi_connection, connection_record): [SO](https://stackoverflow.com/questions/2614984/sqlite-sqlalchemy-how-to-enforce-foreign-keys) """ # Only enable foreign keys if the connection record has the flag set to True - if connection_record.info.get('enable_fks', False): + if connection_record.info.get("enable_fks", False): cursor = dbapi_connection.cursor() cursor.execute("PRAGMA foreign_keys=ON") cursor.close() @@ -80,7 +80,7 @@ def __init__(self, sqlite_path: Optional[str] = None): "isolation_level": "AUTOCOMMIT", # Required for SQLite } self._async_db_engine = create_async_engine(**engine_dict) - self._async_db_engine.connect().connection.connection_record.info['enable_fks'] = True + self._async_db_engine.connect().connection.connection_record.info["enable_fks"] = True def does_db_exist(self): return self._db_path.is_file() From 1dc1da1b849bf55de940660a4d2ae08082c8be20 Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Wed, 22 Jan 2025 10:02:33 +0200 Subject: [PATCH 4/4] Changes to PRAGMA --- migrations/env.py | 1 + src/codegate/db/connection.py | 9 +++------ 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/migrations/env.py b/migrations/env.py index 0729ead0..330a20e0 100644 --- a/migrations/env.py +++ b/migrations/env.py @@ -35,6 +35,7 @@ def run_migrations_offline() -> None: target_metadata=target_metadata, literal_binds=True, dialect_opts={"paramstyle": "named"}, + transactional_ddl=True, ) with context.begin_transaction(): diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index 5f290ffe..dc75bb83 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -45,11 +45,9 @@ def set_sqlite_pragma(dbapi_connection, connection_record): [SQLite docs](https://www.sqlite.org/foreignkeys.html) [SO](https://stackoverflow.com/questions/2614984/sqlite-sqlalchemy-how-to-enforce-foreign-keys) """ - # Only enable foreign keys if the connection record has the flag set to True - if connection_record.info.get("enable_fks", False): - cursor = dbapi_connection.cursor() - cursor.execute("PRAGMA foreign_keys=ON") - cursor.close() + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() class DbCodeGate: @@ -80,7 +78,6 @@ def __init__(self, sqlite_path: Optional[str] = None): "isolation_level": "AUTOCOMMIT", # Required for SQLite } self._async_db_engine = create_async_engine(**engine_dict) - self._async_db_engine.connect().connection.connection_record.info["enable_fks"] = True def does_db_exist(self): return self._db_path.is_file()