From e092c73ec786a8337d1a3da4e7d122b60620159c Mon Sep 17 00:00:00 2001
From: kshitijrajsharma <skshitizraj@gmail.com>
Date: Mon, 25 Dec 2023 11:10:56 +0545
Subject: [PATCH] Set memory limit in duckdb

---
 src/app.py    | 10 +++++++++-
 src/config.py |  4 ++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/app.py b/src/app.py
index c5d70542..055d95fa 100644
--- 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
index f7958bbb..9b841309 100644
--- 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: