diff --git a/nifi/user-scripts/annotation_manager.py b/nifi/user-scripts/annotation_manager.py index 1e4275a6..c629b48e 100644 --- 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 index 6ca078bd..e72eddc1 100644 --- 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: