feature: change sqlite to sqlalchemy+pymysql

Author: syc366

runtime/operators/filter/img_duplicated_images_cleaner/process.py

@@ -17,8 +17,9 @@
 
 import cv2
 from Crypto.Hash import MD5
+from sqlalchemy import text
 
-from data_platform.sqlite_manager.sqlite_manager import SQLiteManager
+from data_platform.sql_manager.sql_manager import SQLManager
 from data_platform.common.utils import get_now_time
 from data_platform.common.utils import bytes_to_numpy, numpy_to_bytes
 from data_platform.core.base_op import Filter
@@ -77,24 +78,24 @@ def execute_sql(self, md5: str, file_name: str,
         query_sql = str(self.sql_dict.get("query_sql"))
         insert_sql = str(self.sql_dict.get("insert_sql"))
         create_tables_sql = str(self.sql_dict.get("create_tables_sql"))
-        db_path = str(self.sql_dict.get("db_path"))
-        query_sql_params = [self.task_uuid, md5]
-        insert_sql_params = [self.task_uuid, md5, file_name.encode("utf-8"), timestamp]
+        query_sql_params = {"task_uuid": self.task_uuid, "file_feature": md5}
+        insert_sql_params = {"task_uuid": self.task_uuid, "file_feature": md5, "file_name": file_name.encode("utf-8"),
+                             "timestamp": timestamp}
 
-        db_manager = SQLiteManager()
+        db_manager = SQLManager()
         try:
-            self.conn = db_manager.create_connect(db_path)
+            self.conn = db_manager.create_connect()
         except Exception as e:
             logger.error("fileName: %s, database connection failed: %s", file_name, str(e))
             raise RuntimeError(82000, str(e)) from None
 
         with self.conn as connection:
-            connection.execute(create_tables_sql)
+            connection.execute(text(create_tables_sql))
             # 判断是否有重复文件
-            result = connection.execute(query_sql, query_sql_params).fetchall()
+            result = connection.execute(text(query_sql, query_sql_params)).fetchall()
             # 查询记录为空，无重复图片, 插入新文件特征
             if not result:
-                connection.execute(insert_sql, insert_sql_params)
+                connection.execute(text(insert_sql, insert_sql_params))
                 return img_bytes
             logger.info("fileName: %s, method: Duplicate ImagesCleaner. The image is duplicated and filtered ",
                         file_name)

runtime/operators/filter/img_duplicated_images_cleaner/sql/sql_config.json

@@ -1,6 +1,5 @@
 {
-  "query_sql": "SELECT * FROM operator_duplicate_img_features WHERE task_uuid = ? AND file_feature = ?",
-  "db_path": "/flow/sqlite.db",
-  "insert_sql": "INSERT INTO operator_duplicate_img_features (task_uuid, file_feature,file_name,timestamp) VALUES (?,?,?,?)",
-  "create_tables_sql": "CREATE TABLE IF NOT EXISTS operator_duplicate_img_features (id INTEGER PRIMARY KEY AUTOINCREMENT,task_uuid TEXT,file_feature TEXT,file_name TEXT,timestamp DATETIME);"
+  "query_sql": "SELECT * FROM operator_duplicate_img_features WHERE task_uuid = :task_uuid AND file_feature = :file_feature",
+  "insert_sql": "INSERT INTO operator_duplicate_img_features (task_uuid, file_feature, file_name, timestamp) VALUES (:task_uuid, :file_feature, :file_name, :timestamp)",
+  "create_tables_sql": "CREATE TABLE IF NOT EXISTS operator_duplicate_img_features (id INT AUTO_INCREMENT PRIMARY KEY,task_uuid VARCHAR(255),file_feature TEXT,file_name TEXT,timestamp DATETIME);"
 }

runtime/operators/filter/img_similar_images_cleaner/process.py

@@ -20,8 +20,9 @@
 
 import cv2
 import numpy as np
+from sqlalchemy import text
 
-from data_platform.sqlite_manager.sqlite_manager import SQLiteManager
+from data_platform.sql_manager.sql_manager import SQLManager
 from data_platform.common.utils import get_now_time
 from data_platform.common.utils import bytes_to_numpy
 from data_platform.core.base_op import Filter
@@ -149,36 +150,48 @@ def execute_sql(self, p_hash: str, des_matrix: np.ndarray, file_name: str,
         des_matrix_binary = zlib.compress(des_matrix.tobytes())  # 使用 zlib 进行压缩数组
         timestamp = get_now_time('Asia/Shanghai', '%Y-%m-%d %H:%M:%S', file_name,
                                  "ImgSimilarCleaner")
-        query_sql = str(self.sql_dict.get("query_sql"))
         query_task_uuid_sql = str(self.sql_dict.get("query_task_uuid_sql"))
         insert_sql = str(self.sql_dict.get("insert_sql"))
         create_tables_sql = str(self.sql_dict.get("create_tables_sql"))
-        db_path = str(self.sql_dict.get("db_path"))
 
-        db_manager = SQLiteManager()
+        db_manager = SQLManager()
         try:
-            self.conn = db_manager.create_connect(db_path)
+            self.conn = db_manager.create_connect()
         except Exception as e:
             logger.error("fileName: %s, database connection failed: %s", file_name, str(e))
             raise RuntimeError(82000, str(e)) from None
 
         with self.conn as connection:
             """从数据库中获取文件特征、比较相似度，插入新的文件特征"""
-            connection.execute(create_tables_sql)
-            result = connection.execute(query_task_uuid_sql, [self.task_uuid]).fetchall()
+            with self.conn as connection:
+                connection.execute(text(create_tables_sql))
+            result = connection.execute(query_task_uuid_sql, {"task_uuid": self.task_uuid}).fetchall()
             total_count = len(result)
-            for i in range(0, total_count, self.page_size):
-                rows = connection.execute(query_sql, [self.task_uuid, self.page_size, i]).fetchall()
-                # 对应任务uuid，最后一页没有数据，跳出循环
-                if not rows:
-                    break
-                # 对两张图片进行相似度比较
-                if self.determine_similar_images(rows, p_hash, des_matrix, file_name):
+            if self.has_similar_images(connection, des_matrix, file_name, p_hash, total_count):
                     return np.array([])
-            connection.execute(insert_sql, [self.task_uuid, p_hash, des_matrix_binary, str(des_matrix.shape),
-                                            file_name.encode("utf-8").hex(), timestamp])
+
+            insert_data = {
+                "task_uuid": self.task_uuid,
+                "p_hash": p_hash,
+                "des_matrix": des_matrix_binary,
+                "matrix_shape": str(des_matrix.shape),
+                "file_name": file_name.encode("utf-8").hex(),
+                "timestamp": timestamp
+            }
+            connection.execute(text(insert_sql),insert_data)
         return img
 
+    def has_similar_images(self, connection, des_matrix, file_name, p_hash, total_count):
+        for i in range(0, total_count, self.page_size):
+            query_sql = self.sql_dict.get("query_sql")
+            rows = connection.execute(text(query_sql), {"task_uuid": self.task_uuid, "ge": self.page_size, "le": i}).fetchall()
+            # 对应任务uuid，最后一页没有数据，跳出循环
+            if not rows:
+                break            # 对两张图片进行相似度比较
+            if self.determine_similar_images(rows, p_hash, des_matrix, file_name):
+                return True
+        return False
+
     def determine_similar_images(self, file_features: List, p_hash: str, des_matrix: np.ndarray,
                                  file_name: str) -> bool:
         """根据文件特征，判断两张图片相似度是否超过指定阈值"""

runtime/operators/filter/img_similar_images_cleaner/sql/sql_config.json

@@ -1,7 +1,6 @@
 {
-  "query_sql": "SELECT * FROM operator_similar_img_features WHERE task_uuid = ? ORDER BY timestamp LIMIT ? OFFSET ?",
-  "insert_sql": "INSERT INTO operator_similar_img_features (task_uuid,p_hash,des_matrix,matrix_shape,file_name,timestamp) VALUES (?,?,?,?,?,?)",
-  "query_task_uuid_sql": "SELECT * FROM operator_similar_img_features WHERE task_uuid = ?",
-  "db_path": "/flow/sqlite.db",
-  "create_tables_sql": "CREATE TABLE IF NOT EXISTS operator_similar_img_features (id INTEGER PRIMARY KEY AUTOINCREMENT,task_uuid TEXT,p_hash TEXT,des_matrix BLOB,matrix_shape TEXT,file_name TEXT,timestamp DATETIME);"
+  "query_sql": "SELECT * FROM operator_similar_img_features WHERE task_uuid = :task_uuid ORDER BY timestamp LIMIT :ge OFFSET :le",
+  "insert_sql": "INSERT INTO operator_similar_img_features (task_uuid,p_hash,des_matrix,matrix_shape,file_name,timestamp) VALUES (:task_uuid,:p_hash,:des_matrix,:matrix_shape,:file_name,:timestamp)",
+  "query_task_uuid_sql": "SELECT * FROM operator_similar_img_features WHERE task_uuid = :task_uuid",
+  "create_tables_sql": "CREATE TABLE IF NOT EXISTS operator_similar_img_features (id INT AUTO_INCREMENT PRIMARY KEY,task_uuid VARCHAR(255),p_hash TEXT,des_matrix BLOB,matrix_shape TEXT,file_name TEXT,timestamp DATETIME);"
 }

runtime/operators/filter/remove_duplicate_file/process.py

@@ -17,8 +17,9 @@
 
 import numpy as np
 from datasketch import MinHash
+from sqlalchemy import text
 
-from data_platform.sqlite_manager.sqlite_manager import SQLiteManager
+from data_platform.sql_manager.sql_manager import SQLManager
 from data_platform.common.utils import get_now_time
 from data_platform.core.base_op import Filter
 
@@ -91,34 +92,43 @@ def execute_sql(self, text_minhash: MinHash, file_name: str,
         minhash_values = text_minhash.hashvalues
         # 将 NumPy 数组转换为字符串
         minhash_values_string = np.array2string(minhash_values)
-
-        query_sql = self.sql_dict.get("query_sql")
         query_task_uuid_sql = self.sql_dict.get("query_task_uuid_sql")
         insert_sql = self.sql_dict.get("insert_sql")
-        db_path = self.sql_dict.get("db_path")
         create_tables_sql = self.sql_dict.get("create_tables_sql")
-        db_manager = SQLiteManager()
+        db_manager = SQLManager()
         try:
-            self.conn = db_manager.create_connect(db_path)
+            self.conn = db_manager.create_connect()
         except Exception as e:
             logger.error("fileName: %s, database connection failed: %s", file_name, str(e))
             raise RuntimeError(82000, str(e)) from None
         with self.conn as connection:
-            connection.execute(create_tables_sql)
-            result = connection.execute(query_task_uuid_sql, [self.task_uuid]).fetchall()
+            connection.execute(text(create_tables_sql))
+            result = connection.execute(query_task_uuid_sql, {"task_uuid": self.task_uuid}).fetchall()
             total_count = len(result)
-            for i in range(0, total_count, self.page_size):
-                rows = connection.execute(query_sql, [self.task_uuid, self.page_size, i]).fetchall()
-                # 对应任务uuid，最后一页没有数据，跳出循环
-                if not rows:
-                    break
-                # 对两个文本进行相似度比较
-                if self.determine_similar_text(rows, text_minhash, file_name):
-                    return ""
-            connection.execute(insert_sql, [self.task_uuid, minhash_values_string, file_name.encode("utf-8").hex(),
-                                            timestamp])
+            if self.has_similar_text(connection, file_name, text_minhash, total_count):
+                return ""
+            insert_data = {
+                "task_uuid": self.task_uuid,
+               "file_feature": minhash_values_string,
+               "file_name": file_name.encode("utf-8").hex(),
+               "timestamp": timestamp
+            }
+            connection.execute(text(insert_sql), insert_data)
         return input_text
 
+    def has_similar_text(self, connection, file_name, text_minhash, total_count) -> bool:
+        query_sql = self.sql_dict.get("query_sql")
+        for i in range(0, total_count, self.page_size):
+            rows = connection.execute(
+                text(query_sql), {"task_uuid": self.task_uuid, "ge": self.page_size, "le": i}).fetchall()
+            # 对应任务uuid，最后一页没有数据，跳出循环
+            if not rows:
+                break
+            # 对两个文本进行相似度比较
+            if self.determine_similar_text(rows, text_minhash, file_name):
+                return True
+        return False
+
     def determine_similar_text(self, file_features: List, text_minhash: MinHash, file_name: str) -> bool:
         for signature in file_features:
             # 历史文件特征和历史文件名称

runtime/operators/filter/remove_duplicate_file/sql/sql_config.json

@@ -1,7 +1,6 @@
 {
-  "query_sql": "SELECT * FROM operators_similar_text_features WHERE task_uuid = ? ORDER BY timestamp LIMIT ? OFFSET ?",
-  "db_path": "/flow/sqlite.db",
-  "create_tables_sql": "CREATE TABLE IF NOT EXISTS operators_similar_text_features (id INTEGER PRIMARY KEY AUTOINCREMENT, task_uuid TEXT,file_feature TEXT,file_name TEXT,timestamp DATETIME);",
-  "insert_sql": "INSERT INTO operators_similar_text_features (task_uuid, file_feature, file_name, timestamp) VALUES (?, ?, ?, ?)",
-  "query_task_uuid_sql": "SELECT * FROM operators_similar_text_features WHERE task_uuid = ?"
+  "query_sql": "SELECT * FROM operators_similar_text_features WHERE task_uuid = :task_uuid ORDER BY timestamp LIMIT :ge OFFSET :le",
+  "create_tables_sql": "CREATE TABLE IF NOT EXISTS operators_similar_text_features (id INT AUTO_INCREMENT PRIMARY KEY, task_uuid VARCHAR(255),file_feature TEXT,file_name TEXT,timestamp DATETIME);",
+  "insert_sql": "INSERT INTO operators_similar_text_features (task_uuid, file_feature, file_name, timestamp) VALUES (:task_uuid, :file_feature, :file_name, :timestamp)",
+  "query_task_uuid_sql": "SELECT * FROM operators_similar_text_features WHERE task_uuid = :task_uuid"
 }

runtime/operators/requirements.txt

@@ -17,3 +17,5 @@ pytz==2025.2
 six==1.17.0
 xmltodict==1.0.2
 zhconv==1.4.3
+sqlalchemy==2.0.40
+pymysql==1.1.1

runtime/python-executor/data_platform/core/base_op.py

@@ -12,7 +12,7 @@
 from data_platform.common.utils.registry import Registry
 from data_platform.common.utils import check_valid_path
 from data_platform.core.constant import Fields
-from data_platform.sqlite_manager.persistence_atction import TaskInfoPersistence
+from data_platform.sql_manager.persistence_atction import TaskInfoPersistence
 
 OPERATORS = Registry('Operators')
 

…data_platform/sqlite_manager/__init__.py …or/data_platform/sql_manager/__init__.pyruntime/python-executor/data_platform/sqlite_manager/__init__.py renamed to runtime/python-executor/data_platform/sql_manager/__init__.py runtime/python-executor/data_platform/sqlite_manager/__init__.py renamed to runtime/python-executor/data_platform/sql_manager/__init__.py

@@ -0,0 +1,138 @@
+# -*- coding: utf-8 -*-
+
+import json
+import time
+import os
+from pathlib import Path
+from typing import Dict, Any
+
+from loguru import logger
+from sqlalchemy import text
+
+from data_platform.sql_manager.sql_manager import SQLManager
+
+
+class TaskInfoPersistence:
+    def __init__(self):
+        self.sql_dict = self.load_sql_dict()
+
+    @staticmethod
+    def load_sql_dict():
+        """获取sql语句"""
+        sql_config_path = str(Path(__file__).parent / 'sql' / 'sql_config.json')
+        with open(sql_config_path, 'r', encoding='utf-8') as f:
+            return json.load(f)
+
+    def persistence_task_info(self, sample: Dict[str, Any]):
+        instance_id = str(sample.get("instance_id"))
+        meta_file_name = str(sample.get("sourceFileName"))
+        meta_file_type = str(sample.get("sourceFileType"))
+        meta_file_id = int(sample.get("sourceFileId"))
+        meta_file_size = str(sample.get("sourceFileSize"))
+        file_id = int(sample.get("fileId"))
+        file_size = str(sample.get("fileSize"))
+        file_type = str(sample.get("fileType"))
+        file_name = str(sample.get("fileName"))
+        file_path = str(sample.get("filePath"))
+        status = int(sample.get("execute_status"))
+        failed_reason = sample.get("failed_reason")
+        operator_id = str(failed_reason.get("op_name")) if failed_reason else ""
+        error_code = str(failed_reason.get("error_code")) if failed_reason else ""
+        incremental = str(sample.get("incremental") if sample.get("incremental") else '')
+        child_id = sample.get("childId")
+        slice_num = sample.get('slice_num', 0)
+        insert_data = {
+            "instance_id": instance_id,
+            "meta_file_id": meta_file_id,
+            "meta_file_type": meta_file_type,
+            "meta_file_name": meta_file_name,
+            "meta_file_size": meta_file_size,
+            "file_id": file_id,
+            "file_size": file_size,
+            "file_type": file_type,
+            "file_name": file_name,
+            "file_path": file_path,
+            "status": status,
+            "operator_id": operator_id,
+            "error_code": error_code,
+            "incremental": incremental,
+            "child_id": child_id,
+            "slice_num": slice_num,
+        }
+        self.insert_clean_result(insert_data)
+
+    def insert_clean_result(self, insert_data):
+        retries = 0
+        max_retries = 20
+        retry_delay = 1
+        while retries <= max_retries:
+            try:
+                self.execute_sql_insert(insert_data)
+                return
+            except Exception as e:
+                if "database is locked" in str(e) or "locking protocol" in str(e):
+                    retries += 1
+                    time.sleep(retry_delay)
+                else:
+                    logger.error("database execute failed: {}", str(e))
+                    raise RuntimeError(82000, str(e)) from None
+        raise Exception("Max retries exceeded")
+
+    def execute_sql_insert(self, insert_data):
+        insert_sql = str(self.sql_dict.get("insert_sql"))
+        create_tables_sql = str(self.sql_dict.get("create_tables_sql"))
+        with SQLManager.create_connect() as conn:
+            conn.execute(text(create_tables_sql))
+            conn.execute(text(insert_sql), insert_data)
+
+
+    def query_task_info(self, instance_ids: list[str]):
+        result = {}
+        current_result = None
+        for instance_id in instance_ids:
+            try:
+                current_result = self.execute_sql_query(instance_id)
+            except Exception as e:
+                logger.warning("instance_id: {}, query job result error: {}", instance_id, str(e))
+            if current_result:
+                result[instance_id] = current_result
+        return result
+
+    def execute_sql_query(self, instance_id):
+        result = None
+        create_tables_sql = str(self.sql_dict.get("create_tables_sql"))
+        query_sql = str(self.sql_dict.get("query_sql"))
+        with SQLManager.create_connect() as conn:
+            conn.execute(text(create_tables_sql))
+            execute_result = conn.execute(text(query_sql), {"instance_id": instance_id})
+            result = execute_result.fetchall()
+        return result
+
+    # todo 删除接口待实现
+    def delete_task_info(self, instance_id: str):
+        create_tables_sql = self.sql_dict.get("create_tables_sql")
+        delete_task_instance_sql = self.sql_dict.get("delete_task_instance_sql")
+        try:
+            with SQLManager.create_connect() as conn:
+                conn.execute(text(create_tables_sql))
+                conn.execute(text(delete_task_instance_sql), {"instance_id": instance_id})
+        except Exception as e:
+            logger.warning(f"delete database for flow: {instance_id}", e)
+
+    def delete_task_operate_info(self, instance_id: str):
+        create_duplicate_img_tables_sql = self.sql_dict.get("create_duplicate_img_tables_sql")
+        create_similar_img_tables_sql = self.sql_dict.get("create_similar_img_tables_sql")
+        create_similar_text_tables_sql = self.sql_dict.get("create_similar_text_tables_sql")
+        delete_duplicate_img_tables_sql = self.sql_dict.get("delete_duplicate_img_tables_sql")
+        delete_similar_img_tables_sql = self.sql_dict.get("delete_similar_img_tables_sql")
+        delete_similar_text_tables_sql = self.sql_dict.get("delete_similar_text_tables_sql")
+        try:
+            with SQLManager.create_connect() as conn:
+                conn.execute(text(create_duplicate_img_tables_sql))
+                conn.execute(text(delete_duplicate_img_tables_sql), {"instance_id": instance_id})
+                conn.execute(text(create_similar_img_tables_sql))
+                conn.execute(text(delete_similar_img_tables_sql), {"instance_id": instance_id})
+                conn.execute(text(create_similar_text_tables_sql))
+                conn.execute(text(delete_similar_text_tables_sql), {"instance_id": instance_id})
+        except Exception as e:
+            logger.warning(f"delete database for flow: {instance_id} error", e)