Set memory limit in duckdb

hotosm · Dec 25, 2023 · e092c73 · e092c73
1 parent 6fd2afd
commit e092c73
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/src/app.py b/src/app.py
@@ -87,7 +87,12 @@
     import duckdb
     from hdx.data.dataset import Dataset
 
-    from src.config import HDX_MAINTAINER, HDX_OWNER_ORG, HDX_URL_PREFIX
+    from src.config import (
+        DUCK_DB_MEMORY_LIMIT,
+        HDX_MAINTAINER,
+        HDX_OWNER_ORG,
+        HDX_URL_PREFIX,
+    )
 
 
 global LOCAL_CON_POOL
@@ -1124,6 +1129,9 @@ def __init__(self, db_path):
         con.install_extension("json")
         con.load_extension("spatial")
         con.load_extension("json")
+        if DUCK_DB_MEMORY_LIMIT:
+            con.sql(f"""SET memory_limit = '{DUCK_DB_MEMORY_LIMIT}'""")
+        con.sql("""SET enable_progress_bar = true""")
 
     def run_query(self, query, attach_pgsql=False, load_spatial=False):
         """

diff --git a/src/config.py b/src/config.py
@@ -228,6 +228,10 @@
         or config.get("HDX", "ALLOWED_HDX_UPDATE_FREQUENCIES", fallback=None)
         or Dataset.list_valid_update_frequencies()
     )
+    DUCK_DB_MEMORY_LIMIT = parse_list(
+        os.environ.get("DUCK_DB_MEMORY_LIMIT")
+        or config.get("HDX", "DUCK_DB_MEMORY_LIMIT", fallback=None)
+    )
 
 
 def get_db_connection_params() -> dict: