1 files changed, 31 insertions, 27 deletions
diff --git a/update_sharing.py b/update_sharing.py
index 55e8096..664b627 100755
--- a/update_sharing.py
+++ b/update_sharing.py
@@ -1,16 +1,18 @@
 #!/usr/bin/python
 
-import sqlite3
+import sqlalchemy
 
-from dedup.utils import fetchiter
+from dedup.utils import fetchiter, enable_sqlite_foreign_keys
 
-def add_values(cursor, insert_key, files, size):
-    cursor.execute("UPDATE sharing SET files = files + ?, size = size + ? WHERE pid1 = ? AND pid2 = ? AND func1 = ? AND func2 = ?;",
-                   (files, size) + insert_key)
-    if cursor.rowcount > 0:
+def add_values(conn, insert_key, files, size):
+    params = dict(files=files, size=size, pid1=insert_key[0],
+                  pid2=insert_key[1], func1=insert_key[2], func2=insert_key[3])
+    rows = conn.execute("UPDATE sharing SET files = files + :files, size = size + :size WHERE pid1 = :pid1 AND pid2 = :pid2 AND func1 = :func1 AND func2 = :func2;",
+                        **params)
+    if rows.rowcount > 0:
         return
-    cursor.execute("INSERT INTO sharing (pid1, pid2, func1, func2, files, size) VALUES (?, ?, ?, ?, ?, ?);",
-                   insert_key + (files, size))
+    conn.execute("INSERT INTO sharing (pid1, pid2, func1, func2, files, size) VALUES (:pid1, :pid2, :func1, :func2, :files, :size);",
+                 **params)
 
 def compute_pkgdict(rows):
     pkgdict = dict()
@@ -19,7 +21,7 @@ def compute_pkgdict(rows):
         funcdict.setdefault(function, []).append((size, filename))
     return pkgdict
 
-def process_pkgdict(cursor, pkgdict):
+def process_pkgdict(conn, pkgdict):
     for pid1, funcdict1 in pkgdict.items():
         for function1, files in funcdict1.items():
             numfiles = len(files)
@@ -35,26 +37,28 @@ def process_pkgdict(cursor, pkgdict):
                     pkgsize = size
                 for function2 in funcdict2.keys():
                     insert_key = (pid1, pid2, function1, function2)
-                    add_values(cursor, insert_key, pkgnumfiles, pkgsize)
+                    add_values(conn, insert_key, pkgnumfiles, pkgsize)
 
 def main():
-    db = sqlite3.connect("test.sqlite3")
-    cur = db.cursor()
-    cur.execute("PRAGMA foreign_keys = ON;")
-    cur.execute("DELETE FROM sharing;")
-    cur.execute("DELETE FROM duplicate;")
-    readcur = db.cursor()
-    readcur.execute("SELECT hash FROM hash GROUP BY hash HAVING count(*) > 1;")
-    for hashvalue, in fetchiter(readcur):
-        cur.execute("SELECT content.pid, content.id, content.filename, content.size, hash.function FROM hash JOIN content ON hash.cid = content.id WHERE hash = ?;",
-                    (hashvalue,))
-        rows = cur.fetchall()
-        print("processing hash %s with %d entries" % (hashvalue, len(rows)))
-        pkgdict = compute_pkgdict(rows)
-        cur.executemany("INSERT OR IGNORE INTO duplicate (cid) VALUES (?);",
-                        [(row[1],) for row in rows])
-        process_pkgdict(cur, pkgdict)
-    db.commit()
+    db = sqlalchemy.create_engine("sqlite:///test.sqlite3")
+    enable_sqlite_foreign_keys(db)
+    with db.begin() as conn:
+        conn.execute("DELETE FROM sharing;")
+        conn.execute("DELETE FROM duplicate;")
+        readcur = conn.execute("SELECT hash FROM hash GROUP BY hash HAVING count(*) > 1;")
+        for hashvalue, in fetchiter(readcur):
+            rows = conn.execute("SELECT content.pid, content.id, content.filename, content.size, hash.function FROM hash JOIN content ON hash.cid = content.id WHERE hash = :hashvalue;",
+                                hashvalue=hashvalue).fetchall()
+            print("processing hash %s with %d entries" % (hashvalue, len(rows)))
+            pkgdict = compute_pkgdict(rows)
+            for row in rows:
+                cid = row[1]
+                already = conn.scalar("SELECT cid FROM duplicate WHERE cid = :cid;",
+                                      cid=cid)
+                if not already:
+                    conn.execute("INSERT INTO duplicate (cid) VALUES (:cid);",
+                                 cid=cid)
+            process_pkgdict(conn, pkgdict)
 
 if __name__ == "__main__":
     main()