feat: pass connection number as params

CAVEconnectome · Apr 12, 2023 · 4ec56a7 · 4ec56a7
1 parent 9c152ec
commit 4ec56a7
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/dynamicannotationdb/database.py b/dynamicannotationdb/database.py
@@ -16,12 +16,12 @@
 
 
 class DynamicAnnotationDB:
-    def __init__(self, sql_url: str) -> None:
+    def __init__(self, sql_url: str, pool_size=5, max_overflow=5) -> None:
 
         self._cached_session = None
         self._cached_tables = {}
         self._engine = create_engine(
-            sql_url, pool_recycle=3600, pool_size=20, max_overflow=50
+            sql_url, pool_recycle=3600, pool_size=pool_size, max_overflow=max_overflow
         )
         self.base = Base
         self.base.metadata.bind = self._engine

diff --git a/dynamicannotationdb/interface.py b/dynamicannotationdb/interface.py
@@ -35,11 +35,15 @@ class DynamicAnnotationInterface:
 
     """
 
-    def __init__(self, url: str, aligned_volume: str) -> None:
+    def __init__(
+        self, url: str, aligned_volume: str, pool_size=5, max_overflow=5
+    ) -> None:
         self._annotation = None
         self._database = None
         self._segmentation = None
         self._schema = None
+        self.pool_size = pool_size
+        self.max_overflow = max_overflow
         self._base_url = url.rpartition("/")[0]
         self._aligned_volume = aligned_volume
         self._sql_url = self.create_or_select_database(url, aligned_volume)
@@ -159,7 +163,9 @@ def annotation(self) -> DynamicAnnotationClient:
     @property
     def database(self) -> DynamicAnnotationDB:
         if not self._database:
-            self._database = DynamicAnnotationDB(self._sql_url)
+            self._database = DynamicAnnotationDB(
+                self._sql_url, self.pool_size, self.max_overflow
+            )
         return self._database
 
     @property