enable result buffering for postgres

4 files changed, 16 insertions, 6 deletions

diff --git a/autoimport.py b/autoimport.py
index 6db6171..7d5c000 100755
--- a/autoimport.py
+++ b/autoimport.py
@@ -18,7 +18,7 @@ from debian import deb822
 from debian.debian_support import version_compare
 import sqlalchemy
 
-from dedup.utils import enable_sqlite_foreign_keys
+from dedup.utils import configure_database_engine
 
 from readyaml import readyaml
 
@@ -91,7 +91,7 @@ def main():
     options, args = parser.parse_args()
     tmpdir = tempfile.mkdtemp(prefix=b"debian-dedup")
     db = sqlalchemy.create_engine(options.database)
-    enable_sqlite_foreign_keys(db)
+    configure_database_engine(db)
     e = concurrent.futures.ThreadPoolExecutor(multiprocessing.cpu_count())
     pkgs = {}
     for d in args:
diff --git a/dedup/utils.py b/dedup/utils.py
index e163a00..1452edc 100644
--- a/dedup/utils.py
+++ b/dedup/utils.py
@@ -22,3 +22,13 @@ def sqlite_add_version_compare(engine):
     def add_version_compare(connection, _):
         connection.create_collation("debian_version", version_compare)
         connection.create_function("debian_version_compare", 2, version_compare)
+
+def postgres_enable_result_streaming(engine):
+    if engine.dialect.driver != "psycopg2":
+        return
+    engine.server_side_cursors = True
+
+def configure_database_engine(engine):
+    enable_sqlite_foreign_keys(engine)
+    sqlite_add_version_compare(engine)
+    postgres_enable_result_streaming(engine)
diff --git a/readyaml.py b/readyaml.py
index 15cfcb3..8ba82fa 100755
--- a/readyaml.py
+++ b/readyaml.py
@@ -9,7 +9,7 @@ from debian.debian_support import version_compare
 import sqlalchemy
 import yaml
 
-from dedup.utils import enable_sqlite_foreign_keys
+from dedup.utils import configure_database_engine
 
 def readyaml(conn, stream):
     gen = yaml.safe_load_all(stream)
@@ -64,7 +64,7 @@ def main():
                       help="location of the database")
     options, args = parser.parse_args()
     db = sqlalchemy.create_engine(options.database)
-    enable_sqlite_foreign_keys(db)
+    configure_database_engine(db)
     with db.begin() as conn:
         readyaml(conn, sys.stdin)
 
diff --git a/update_sharing.py b/update_sharing.py
index 450bfc7..e41bceb 100755
--- a/update_sharing.py
+++ b/update_sharing.py
@@ -4,7 +4,7 @@ import optparse
 
 import sqlalchemy
 
-from dedup.utils import fetchiter, enable_sqlite_foreign_keys
+from dedup.utils import fetchiter, configure_database_engine
 
 def add_values(conn, insert_key, files, size):
     params = dict(files=files, size=size, pid1=insert_key[0],
@@ -78,5 +78,5 @@ if __name__ == "__main__":
                       help="location of the database")
     options, args = parser.parse_args()
     db = sqlalchemy.create_engine(options.database)
-    enable_sqlite_foreign_keys(db)
+    configure_database_engine(db)
     main(db)