Added pragma on file connections and DROP TABLE IF EXISTS

Author: aponcedeleonch

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);")

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()