dedup.utils: add enbale_sqlite_foreign_keys helper

Makes usage of sqlalchemy easier, cause I can invoke it once and it
works for all connections.

2 files changed, 9 insertions, 2 deletions

diff --git a/dedup/utils.py b/dedup/utils.py
index 2fae9fd..e4d1c10 100644
--- a/dedup/utils.py
+++ b/dedup/utils.py
@@ -1,3 +1,5 @@
+import sqlalchemy.event
+
 def fetchiter(cursor):
     rows = cursor.fetchmany()
     while rows:
@@ -5,3 +7,8 @@ def fetchiter(cursor):
             yield row
         rows = cursor.fetchmany()
 
+def enable_sqlite_foreign_keys(engine):
+    @sqlalchemy.event.listens_for(engine, "connect")
+    def pragma_foreign_keys(connection, _):
+        connection.execute("PRAGMA foreign_keys=ON;")
+
diff --git a/update_sharing.py b/update_sharing.py
index bbd19e5..e78e13f 100755
--- a/update_sharing.py
+++ b/update_sharing.py
@@ -2,7 +2,7 @@
 
 import sqlalchemy
 
-from dedup.utils import fetchiter
+from dedup.utils import fetchiter, enable_sqlite_foreign_keys
 
 def add_values(conn, insert_key, files, size):
     rows = conn.execute("UPDATE sharing SET files = files + ?, size = size + ? WHERE package1 = ? AND package2 = ? AND func1 = ? AND func2 = ?;",
@@ -39,8 +39,8 @@ def process_pkgdict(conn, pkgdict):
 
 def main():
     db = sqlalchemy.create_engine("sqlite:///test.sqlite3")
+    enable_sqlite_foreign_keys(db)
     with db.begin() as conn:
-        conn.execute("PRAGMA foreign_keys = ON;")
         conn.execute("DELETE FROM sharing;")
         conn.execute("DELETE FROM duplicate;")
         readcur = conn.execute("SELECT hash FROM hash GROUP BY hash HAVING count(*) > 1;")