Updated NiFi scripts.

CogStack · Dec 17, 2024 · 223a136 · 223a136
1 parent 0dcf0d8
commit 223a136
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/nifi/user-scripts/annotation_manager.py b/nifi/user-scripts/annotation_manager.py
@@ -83,7 +83,7 @@ def main():
         for record in records:
             if OPERATION_MODE == "check":
                 document_id = str(record[DOCUMENT_ID_FIELD_NAME])
-                query = "SELECT id FROM annotations WHERE id LIKE '%" + document_id + "%' LIMIT 1"
+                query = "SELECT id FROM annotations WHERE document_id = " + document_id + " LIMIT 1"
                 result = connect_and_query(query, db_file_path, sqlite_connection=_sqlite_connection_ro, cursor=_cursor, keep_conn_open=True)
 
                 if len(result) < 1:
@@ -92,7 +92,7 @@ def main():
             if OPERATION_MODE == "insert":
                 document_id = str(record["meta." + DOCUMENT_ID_FIELD_NAME])
                 nlp_id = str(record["nlp.id"])
-                query = "INSERT OR REPLACE INTO annotations (id) VALUES (" + '"' + document_id + "_" + nlp_id + '"' + ")"
+                query = "INSERT OR REPLACE INTO annotations (id, document_id) VALUES (" + '"' + document_id + "_" + nlp_id + '"' + "," + '"' + document_id + '"' +")"
                 result = connect_and_query(query, db_file_path, sqlite_connection=_sqlite_connection_rw, sql_script_mode=True, cursor=_cursor, keep_conn_open=True)
                 output_stream.append(record)
 

diff --git a/nifi/user-scripts/utils/sqlite_query.py b/nifi/user-scripts/utils/sqlite_query.py
@@ -52,9 +52,13 @@ def create_connection(db_file_path: str, read_only_mode=False) -> sqlite3.Connec
         connection_str += "?mode=ro"
 
     _tmp_conn = sqlite3.connect(connection_str, uri=True)
-    _tmp_conn.execute("pragma journal_mode = wal")
-    _tmp_conn.execute("pragma synchronous = normal")
-    _tmp_conn.execute("pragma journal_size_limit = 6144000")
+
+    if read_only_mode:
+        _tmp_conn.execute("pragma synchronous = OFF")
+    else:
+        _tmp_conn.execute("pragma journal_mode = wal")
+        _tmp_conn.execute("pragma synchronous = normal")
+        _tmp_conn.execute("pragma journal_size_limit = 6144000")
 
     return _tmp_conn