Ann manager sql rw conn wal mode.

CogStack · Dec 13, 2024 · e8175a7 · e8175a7
1 parent 3847119
commit e8175a7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/nifi/user-scripts/annotation_manager.py b/nifi/user-scripts/annotation_manager.py
@@ -59,6 +59,7 @@ def main():
         records = json_data_records
 
         _sqlite_connection_ro = None
+        _sqlite_connection_rw = None
 
         if isinstance(records, dict):
             if "content" in json_data_records.keys():
@@ -75,6 +76,7 @@ def main():
         if OPERATION_MODE == "insert":
             del output_stream
             output_stream = []
+            _sqlite_connection_rw = create_connection(db_file_path, read_only_mode=False)
 
         for record in records:
             if OPERATION_MODE == "check":
@@ -89,7 +91,7 @@ def main():
                 document_id = str(record["meta." + DOCUMENT_ID_FIELD_NAME])
                 nlp_id = str(record["nlp.id"])
                 query = "INSERT OR REPLACE INTO annotations (elasticsearch_id) VALUES (" + '"' + document_id + "_" + nlp_id + '"' + ")"
-                result = connect_and_query(query, db_file_path, sql_script_mode=True)
+                result = connect_and_query(query, db_file_path, sqlite_connection=_sqlite_connection_rw, sql_script_mode=True, keep_conn_open=True)
                 output_stream.append(record)
 
         if _sqlite_connection_ro is not None:

diff --git a/nifi/user-scripts/utils/sqlite_query.py b/nifi/user-scripts/utils/sqlite_query.py
@@ -24,7 +24,6 @@ def connect_and_query(query: str, db_file_path: str, sqlite_connection: sqlite3.
             sqlite_connection = sqlite_connection
         else:
             sqlite_connection = create_connection(db_file_path)
-            sqlite_connection.execute('pragma journal_mode=wal')
 
         cursor = sqlite_connection.cursor()
         if not sql_script_mode:
@@ -50,7 +49,10 @@ def create_connection(db_file_path: str, read_only_mode=False) -> sqlite3.Connec
     if read_only_mode:
         connection_str += "?mode=ro"
 
-    return sqlite3.connect(connection_str, uri=True)
+    _tmp_conn = sqlite3.connect(connection_str, uri=True)
+    _tmp_conn.execute('pragma journal_mode=wal')
+
+    return _tmp_conn
 
 
 def query_with_connection(query: str, sqlite_connection: sqlite3.Connection) -> List: