'commit'

DataX/tool/LoadCsvToMysql.py

@@ -0,0 +1,143 @@
+import pymysql
+import sys
+import os
+import glob
+import json
+import time
+from urllib.parse import urlparse, parse_qs
+
+def parse_jdbc_url(url):
+    # jdbc:mysql://host:port/db?params
+    # remove jdbc:mysql://
+    if url.startswith("jdbc:mysql://"):
+        url = url[13:]
+    
+    # split host:port and db
+    if "/" in url:
+        address, remainder = url.split("/", 1)
+        if "?" in remainder:
+            db, params = remainder.split("?", 1)
+        else:
+            db = remainder
+            params = ""
+    else:
+        address = url
+        db = ""
+        params = ""
+        
+    if ":" in address:
+        host, port = address.split(":")
+        port = int(port)
+    else:
+        host = address
+        port = 3306
+        
+    return host, port, db
+
+def load_csv(jdbc_url, user, password, table, csv_dir, columns=None):
+    host, port, db = parse_jdbc_url(jdbc_url)
+    
+    print(f"Connecting to MySQL {host}:{port}/{db} as {user}...")
+    
+    try:
+        conn = pymysql.connect(
+            host=host,
+            port=port,
+            user=user,
+            password=password,
+            database=db,
+            local_infile=True,
+            charset='utf8mb4',
+            cursorclass=pymysql.cursors.DictCursor
+        )
+    except Exception as e:
+        print(f"Connection failed: {e}")
+        sys.exit(1)
+        
+    try:
+        with conn.cursor() as cursor:
+            # Optimization settings
+            print("Setting session parameters for speed...")
+            cursor.execute("SET NAMES utf8mb4")
+            try:
+                cursor.execute("SET FOREIGN_KEY_CHECKS = 0")
+                cursor.execute("SET UNIQUE_CHECKS = 0")
+                cursor.execute("SET SQL_LOG_BIN = 0")
+            except Exception as e:
+                print(f"Warning: Could not set some optimization flags: {e}")
+
+            # Truncate table
+            print(f"Truncating table {table}...")
+            cursor.execute(f"TRUNCATE TABLE `{table}`")
+            
+            # Find files
+            files = glob.glob(os.path.join(csv_dir, "*"))
+            if not files:
+                print(f"No files found in {csv_dir}")
+                return
+
+            total_rows = 0
+            start_time = time.time()
+            
+            for file_path in files:
+                file_path = os.path.abspath(file_path).replace('\\', '/')
+                print(f"Loading file: {file_path}")
+                
+                # Build SQL
+                # Assuming DataX txtfilewriter defaults:
+                # separator: ,
+                # quoteChar: "
+                # escapeChar: \
+                # nullFormat: \N
+                
+                col_sql = ""
+                if columns:
+                    col_list = [f"`{c}`" for c in columns]
+                    col_sql = "(" + ", ".join(col_list) + ")"
+                
+                sql = f"""
+                    LOAD DATA LOCAL INFILE '{file_path}'
+                    INTO TABLE `{table}`
+                    CHARACTER SET utf8mb4
+                    FIELDS TERMINATED BY ',' 
+                    OPTIONALLY ENCLOSED BY '"' 
+                    ESCAPED BY '\\\\'
+                    LINES TERMINATED BY '\\n'
+                    {col_sql}
+                """
+                
+                cursor.execute(sql)
+                rows = cursor.rowcount
+                total_rows += rows
+                print(f"  -> Loaded {rows} rows")
+                
+            conn.commit()
+            
+            duration = time.time() - start_time
+            print(f"Total loaded: {total_rows} rows in {duration:.2f}s ({total_rows/duration if duration > 0 else 0:.2f} rows/s)")
+            
+    except Exception as e:
+        print(f"Error during load: {e}")
+        sys.exit(1)
+    finally:
+        conn.close()
+
+if __name__ == "__main__":
+    if len(sys.argv) < 6:
+        print("Usage: python LoadCsvToMysql.py <jdbc_url> <user> <password> <table> <csv_dir> [columns_json]")
+        sys.exit(1)
+        
+    jdbc_url = sys.argv[1]
+    user = sys.argv[2]
+    password = sys.argv[3]
+    table = sys.argv[4]
+    csv_dir = sys.argv[5]
+    
+    columns = None
+    if len(sys.argv) > 6:
+        try:
+            columns = json.loads(sys.argv[6])
+        except:
+            print("Warning: Could not parse columns JSON")
+            
+    load_csv(jdbc_url, user, password, table, csv_dir, columns)