separate locks

datahub-project · Nov 19, 2021 · 6d2df18 · 6d2df18
1 parent 25887dd
commit 6d2df18
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py b/metadata-ingestion/src/datahub/utilities/sqlalchemy_query_combiner.py
@@ -119,12 +119,15 @@ class SQLAlchemyQueryCombiner:
     # There will be one main greenlet per thread. As such, queries will be
     # queued according to the main greenlet's thread ID. We also keep track
     # of the greenlets we spawn for bookkeeping purposes.
-    _thread_unsafe_operation_lock: threading.Lock = dataclasses.field(
+    _queries_by_thread_lock: threading.Lock = dataclasses.field(
+        default_factory=lambda: threading.Lock()
+    )
+    _greenlets_by_thread_lock: threading.Lock = dataclasses.field(
         default_factory=lambda: threading.Lock()
     )
     _queries_by_thread: Dict[
         greenlet.greenlet, Dict[str, _QueryFuture]
-    ] = dataclasses.field(default_factory=dict)
+    ] = dataclasses.field(default_factory=lambda: collections.defaultdict(dict))
     _greenlets_by_thread: Dict[
         greenlet.greenlet, Set[greenlet.greenlet]
     ] = dataclasses.field(default_factory=lambda: collections.defaultdict(set))
@@ -144,15 +147,15 @@ def _get_main_greenlet(self) -> greenlet.greenlet:
     def _get_queue(self, main_greenlet: greenlet.greenlet) -> Dict[str, _QueryFuture]:
         assert main_greenlet.parent is None
 
-        with self._thread_unsafe_operation_lock:
+        with self._queries_by_thread_lock:
             return self._queries_by_thread.setdefault(main_greenlet, {})
 
     def _get_greenlet_pool(
         self, main_greenlet: greenlet.greenlet
     ) -> Set[greenlet.greenlet]:
         assert main_greenlet.parent is None
 
-        with self._thread_unsafe_operation_lock:
+        with self._greenlets_by_thread_lock:
             return self._greenlets_by_thread[main_greenlet]
 
     def _handle_execute(