[proxy] guard more code with lock in Haketilo state implementation

2 files changed, 30 insertions, 25 deletions

diff --git a/src/hydrilla/proxy/state_impl/base.py b/src/hydrilla/proxy/state_impl/base.py
index 75d733f..f969b19 100644
--- a/src/hydrilla/proxy/state_impl/base.py
+++ b/src/hydrilla/proxy/state_impl/base.py
@@ -153,10 +153,10 @@ class HaketiloStateWithFields(st.HaketiloState):
     @contextmanager
     def cursor(self, transaction: bool = False) \
         -> t.Iterator[sqlite3.Cursor]:
-        """...."""
-        start_transaction = transaction and not self.connection.in_transaction
-
         with self.lock:
+            start_transaction = \
+                transaction and not self.connection.in_transaction
+
             if self.current_cursor is not None:
                 yield self.current_cursor
                 return
diff --git a/src/hydrilla/proxy/state_impl/concrete_state.py b/src/hydrilla/proxy/state_impl/concrete_state.py
index 3611db1..cd32e83 100644
--- a/src/hydrilla/proxy/state_impl/concrete_state.py
+++ b/src/hydrilla/proxy/state_impl/concrete_state.py
@@ -171,28 +171,7 @@ class ConcreteHaketiloState(base.HaketiloStateWithFields):
 
             _operations.pull_missing_files(cursor)
 
-    def rebuild_structures(self) -> None:
-        with self.cursor() as cursor:
-            cursor.execute(
-                '''
-                SELECT
-                        p.payload_id,
-                        p.pattern,
-                        p.eval_allowed,
-                        p.cors_bypass_allowed,
-                        ms.enabled,
-                        i.identifier
-                FROM
-                             payloads         AS p
-                        JOIN item_versions    AS iv
-                                ON p.mapping_item_id = iv.item_version_id
-                        JOIN items            AS i  USING (item_id)
-                        JOIN mapping_statuses AS ms USING (item_id);
-                '''
-            )
-
-            rows = cursor.fetchall()
-
+    def _rebuild_structures(self, cursor: sqlite3.Cursor) -> None:
         new_policy_tree = base.PolicyTree()
 
         ui_factory = policies.WebUIPolicyFactory(builtin=True)
@@ -203,6 +182,28 @@ class ConcreteHaketiloState(base.HaketiloStateWithFields):
                 ui_factory
             )
 
+        cursor.execute(
+            '''
+            SELECT
+                    p.payload_id,
+                    p.pattern,
+                    p.eval_allowed,
+                    p.cors_bypass_allowed,
+                    ms.enabled,
+                    i.identifier
+            FROM
+                         payloads         AS p
+                    JOIN item_versions    AS iv
+                            ON p.mapping_item_id = iv.item_version_id
+                    JOIN items            AS i
+                            USING (item_id)
+                    JOIN mapping_statuses AS ms
+                            USING (item_id);
+            '''
+        )
+
+        rows = cursor.fetchall()
+
         new_payloads_data: dict[st.PayloadRef, st.PayloadData] = {}
 
         for row in rows:
@@ -244,6 +245,10 @@ class ConcreteHaketiloState(base.HaketiloStateWithFields):
         self.policy_tree   = new_policy_tree
         self.payloads_data = new_payloads_data
 
+    def rebuild_structures(self) -> None:
+        with self.cursor() as cursor:
+            self._rebuild_structures(cursor)
+
     def repo_store(self) -> st.RepoStore:
         return repos.ConcreteRepoStore(self)