Let resource limits (threads + memory) be overridable

Author: jonasdedden

smoosense-py/smoosense/utils/duckdb_connections.py

@@ -16,9 +16,17 @@
 
 @validate_call()
 def duckdb_connection_default() -> DuckdbConnectionMaker:
+    # Optional caps via env vars; if unset, DuckDB picks defaults from the cgroup / nproc.
+    memory_limit = os.getenv("SMOOSENSE_DUCKDB_MEMORY_LIMIT")
+    threads = os.getenv("SMOOSENSE_DUCKDB_THREADS")
+
     def maker() -> DuckDBPyConnection:
         con = duckdb.connect()
 
+        if memory_limit:
+            con.execute(f"SET memory_limit='{memory_limit}'")
+        if threads:
+            con.execute(f"SET threads={threads}")
         # Now install and load the extension
         con.execute("INSTALL httpfs")
         con.execute("LOAD httpfs")
@@ -32,7 +40,6 @@ def maker() -> DuckDBPyConnection:
 @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
 def duckdb_connection_using_s3(
     s3_client: BaseClient | None = None,
-    memory_limit: str = "3GB",
 ) -> DuckdbConnectionMaker:
     if s3_client is None:
         s3_client = boto3.client("s3")
@@ -62,7 +69,6 @@ def maker() -> DuckDBPyConnection:
         # Set home_directory and temp_directory before httpfs auto-installs
         con.execute(f"SET home_directory='{home_directory}'")
         con.execute(f"SET temp_directory='{temp_directory}'")
-        con.execute(f"SET memory_limit='{memory_limit}'")
 
         # Now install and load the extension
         con.execute("INSTALL httpfs")