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: