import os import mysql.connector from mysql.connector import Error as MySQLError import json import argparse import sys # Database operations def connect_to_database(tidb_host, tidb_port, database): print("Connecting to tidb, host={} port={} database={}".format(tidb_host, tidb_port, database)) try: connection = mysql.connector.connect( host=tidb_host, user="root", password="", database=database, port=tidb_port, connection_timeout=300, ) if connection.is_connected(): print("Successfully connected to the database") return connection except MySQLError as e: print(f"Error connecting to the database: {e}") return None def read_hdfs_logs(file_path, max_rows=None, batch_size=50000): """Read HDFS logs from a JSON file with optional row limit Returns a generator that yields batches of logs to avoid loading all into memory""" total_read = 0 try: with open(file_path, 'r') as file: batch = [] for i, line in enumerate(file): if max_rows is not None and i >= max_rows: break try: log_entry = json.loads(line.strip()) batch.append(log_entry) total_read += 1 # Yield batch when it reaches the batch size if len(batch) >= batch_size: yield batch batch = [] except json.JSONDecodeError as e: print(f"Error parsing JSON at line {i+1}: {e}") # Yield the remaining batch if any if batch: yield batch except Exception as e: print(f"Error reading file {file_path}: {e}") print(f"Read {total_read} total log entries from {file_path}") def create_hdfs_log_table(connection, table_name): """Create a table for HDFS logs with tenant_id encoded in primary key""" cursor = None try: cursor = connection.cursor() cursor.execute(f""" DROP TABLE IF EXISTS {table_name}; """) print(f"Table {table_name} dropped successfully") cursor.execute(f""" CREATE TABLE IF NOT EXISTS {table_name} ( id BIGINT AUTO_INCREMENT, timestamp BIGINT, severity_text VARCHAR(50), body TEXT, tenant_id INT, PRIMARY KEY (tenant_id, id) )AUTO_INCREMENT = 1000; """) print(f"Table {table_name} created successfully") except MySQLError as e: print(f"Error creating table: {e}") finally: if cursor: cursor.close() def insert_hdfs_logs_batch(connection, table_name, logs): """Insert HDFS logs into database with tenant_id encoded in primary key""" cursor = None try: # Check if connection is still valid if not connection.is_connected(): print("Connection lost, attempting to reconnect...") connection.reconnect(attempts=3, delay=2) cursor = connection.cursor() total_logs = len(logs) sql = f"INSERT INTO {table_name} (timestamp, severity_text, body, tenant_id) VALUES (%s, %s, %s, %s)" batch_data = [] for log in logs: # Extract the required fields timestamp = log.get('timestamp') severity_text = log.get('severity_text') body = log.get('body') tenant_id = log.get('tenant_id') batch_data.append((timestamp, severity_text, body, tenant_id)) cursor.executemany(sql, batch_data) connection.commit() print(f"{total_logs} logs from batch inserted successfully into {table_name}") finally: if cursor: cursor.close() def write_logs_to_csv(logs, outfile): """Write logs to a CSV file""" import csv try: fieldnames = ['timestamp', 'severity_text', 'body', 'tenant_id'] writer = csv.DictWriter(outfile, fieldnames=fieldnames) for log in logs: writer.writerow({ 'timestamp': log.get('timestamp'), 'severity_text': log.get('severity_text'), 'body': log.get('body'), 'tenant_id': log.get('tenant_id') }) outfile.flush() print(f"Logs written to file successfully") except Exception as e: print(f"Error writing logs to CSV: {e}") def process_hdfs_logs( table_name, input_file=None, max_rows=None, batch_size=50000, tidb_host="localhost", tidb_port=4000, database="test", out=None, create_table=True, ): """Process HDFS logs and insert them into the database in batches""" success = True conn = None outfile = None if out: if not out.endswith('.csv'): raise ValueError( "Output file must end with .csv if --out is specified") else: outfile = open(out, 'w', newline='') else: # Connect to tidb conn = connect_to_database(tidb_host, tidb_port, database) if not conn: print("Failed to connect to the database") return False if create_table: create_hdfs_log_table(conn, table_name) else: print(f"Skipping table creation, inserting into existing table {table_name}") try: total_inserted = 0 if input_file is None: asset_dir = os.getenv('ASSET_DIR', "").rstrip('/') print( f"Loading input file from ASSET_DIR env, ASSET_DIR={asset_dir}") infilename = '{}/hdfs-logs-multitenants.json'.format(asset_dir) else: infilename = input_file print( f"Processing logs from '{infilename}' in batches of {batch_size}") # Process logs in batches using the generator for _, batch in enumerate(read_hdfs_logs(infilename, max_rows, batch_size)): if out: write_logs_to_csv(batch, outfile) else: insert_hdfs_logs_batch(conn, table_name, batch) total_inserted += len(batch) print(f"Total logs inserted so far: {total_inserted}") except MySQLError as e: print(f"Database error processing HDFS logs: {e}") success = False except Exception as e: print(f"Error processing HDFS logs: {e}") success = False finally: if conn: conn.close() if outfile: outfile.close() return success if __name__ == "__main__": parser = argparse.ArgumentParser( description='Process HDFS logs and insert into database') parser.add_argument('table_name', type=str, help='Name of the database table') parser.add_argument('--max_rows', type=int, default=None, help='Maximum number of rows to process (optional)') parser.add_argument('--batch_size', type=int, default=50000, help='Number of rows to process in each batch (default: 50000)') parser.add_argument('--tidb_host', type=str, default='localhost', help='TiDB address to connect to (default: localhost)') parser.add_argument('--tidb_port', type=int, default=4000, help='TiDB port to connect to (default: 4000)') parser.add_argument('--database', type=str, default='test', help='Database name to connect to (default: test)') parser.add_argument('--input', type=str, help='Path to input file', default=None) parser.add_argument('--out', type=str, default=None, help='Output file for CSV mode') parser.add_argument( '--insert_only', action='store_true', help='Insert into existing table only, do not create/drop table', ) args = parser.parse_args() print(f"\nProcessing HDFS logs into database table '{args.table_name}':") success = process_hdfs_logs( args.table_name, args.input, args.max_rows, args.batch_size, args.tidb_host, args.tidb_port, args.database, args.out, create_table=not args.insert_only, ) print("\nData processing complete." if success else "\nData processing failed.") sys.exit(0 if success else 1)