Browse Source

Update db_helper: add batch_insert to MySQLHelper

StrayWarrior 7 tháng trước cách đây
mục cha
commit
5246ac92de
1 tập tin đã thay đổi với 26 bổ sung0 xóa
  1. 26 0
      db_helper.py

+ 26 - 0
db_helper.py

@@ -362,6 +362,32 @@ class MysqlHelper(object):
         conn.close()
         return data
 
+    def batch_insert(self, table, data, columns=None):
+        """
+        data: data, list[tuple] or list[dict]
+        columns: column names, list, required if data is list[tuple]
+        """
+        if data is None or len(data) == 0:
+            return
+        conn = pymysql.connect(**self.mysql_info)
+        try:
+            if columns is not None:
+                if len(data[0]) != len(columns):
+                    raise Exception("data length != column length")
+                columns_str = ','.join(columns)
+            else:
+                if isinstance(data[0], dict):
+                    columns_str = ','.join(data[0].keys())
+            placeholders_str = ','.join(['%s'] * len(data[0]))
+            with conn.cursor() as cursor:
+                sql_str = f"INSERT INTO {table} ({columns_str}) VALUES ({placeholders_str})"
+                cursor.executemany(sql_str, data)
+                conn.commit()
+        except pymysql.MySQLError as e:
+            print(f"Error in batch_insert: {e}")
+            connection.rollback()
+        conn.close()
+
 
 if __name__ == '__main__':
     redis_helper = RedisHelper()