798 lines
42 KiB
Python
798 lines
42 KiB
Python
"""
|
||
MySQL to ChromaDB synchronization module
|
||
"""
|
||
import urllib
|
||
import pymysql
|
||
import requests
|
||
from typing import List, Dict, Optional, Tuple
|
||
from datetime import datetime
|
||
from loguru import logger
|
||
from config import DatabaseConfig, settings
|
||
from rag.file_parser import FileParser
|
||
|
||
|
||
class MySQLSync:
|
||
"""Handle synchronization between MySQL and ChromaDB"""
|
||
|
||
def __init__(self, db_config: DatabaseConfig, connection=None):
|
||
"""
|
||
Initialize MySQL sync with database configuration
|
||
|
||
Args:
|
||
db_config: DatabaseConfig object containing database table info
|
||
connection: Optional existing MySQL connection (for multi-database scenario)
|
||
"""
|
||
self.file_parser = FileParser()
|
||
self.db_config = db_config
|
||
self.connection = connection
|
||
if self.connection is None:
|
||
self._connect()
|
||
else:
|
||
# Use existing connection, just switch database
|
||
self._switch_database()
|
||
|
||
def _connect(self):
|
||
"""Establish MySQL connection using config-specific or default connection info"""
|
||
try:
|
||
# Use database-specific connection info if provided, otherwise use .env defaults
|
||
host = self.db_config.mysql_host if self.db_config.mysql_host is not None else settings.MYSQL_HOST
|
||
port = self.db_config.mysql_port if self.db_config.mysql_port is not None else settings.MYSQL_PORT
|
||
user = self.db_config.mysql_user if self.db_config.mysql_user is not None else settings.MYSQL_USER
|
||
password = self.db_config.mysql_password if self.db_config.mysql_password is not None else settings.MYSQL_PASSWORD
|
||
charset = self.db_config.mysql_charset if self.db_config.mysql_charset is not None else settings.MYSQL_CHARSET
|
||
|
||
self.connection = pymysql.connect(
|
||
host=host,
|
||
port=port,
|
||
user=user,
|
||
password=password,
|
||
database=self.db_config.database,
|
||
charset=charset,
|
||
cursorclass=pymysql.cursors.DictCursor
|
||
)
|
||
connection_info = f"{host}:{port}/{self.db_config.database}"
|
||
logger.info(f"Connected to MySQL database: {connection_info} (name: {self.db_config.name})")
|
||
except Exception as e:
|
||
logger.error(f"Failed to connect to MySQL database {self.db_config.name}: {e}")
|
||
raise
|
||
|
||
@staticmethod
|
||
def check_mysql_connection() -> Tuple[bool, str]:
|
||
"""
|
||
Check if MySQL server is accessible and connection can be established
|
||
|
||
Returns:
|
||
Tuple of (is_connected: bool, error_message: str)
|
||
If connected, error_message will be empty string
|
||
"""
|
||
temp_connection = None
|
||
try:
|
||
logger.info(f"Checking MySQL connection to {settings.MYSQL_HOST}:{settings.MYSQL_PORT}...")
|
||
temp_connection = pymysql.connect(
|
||
host=settings.MYSQL_HOST,
|
||
port=settings.MYSQL_PORT,
|
||
user=settings.MYSQL_USER,
|
||
password=settings.MYSQL_PASSWORD,
|
||
charset=settings.MYSQL_CHARSET,
|
||
cursorclass=pymysql.cursors.DictCursor,
|
||
connect_timeout=10 # 10 seconds timeout
|
||
)
|
||
|
||
# Test the connection by executing a simple query
|
||
with temp_connection.cursor() as cursor:
|
||
cursor.execute("SELECT VERSION() as version")
|
||
result = cursor.fetchone()
|
||
mysql_version = result.get('version', 'unknown') if result else 'unknown'
|
||
|
||
logger.info(f"✓ MySQL connection successful (version: {mysql_version})")
|
||
return True, ""
|
||
except pymysql.err.OperationalError as e:
|
||
# Safely get error code and message
|
||
error_code = e.args[0] if len(e.args) > 0 else None
|
||
error_msg = e.args[1] if len(e.args) > 1 else str(e)
|
||
if error_code == 2003:
|
||
error_message = f"Cannot connect to MySQL server at {settings.MYSQL_HOST}:{settings.MYSQL_PORT}. " \
|
||
f"Please check if MySQL server is running and accessible."
|
||
elif error_code == 1045:
|
||
error_message = f"Access denied for user '{settings.MYSQL_USER}'@{settings.MYSQL_HOST}. " \
|
||
f"Please check MySQL username and password."
|
||
else:
|
||
error_message = f"MySQL connection error ({error_code}): {error_msg}"
|
||
logger.error(f"✗ MySQL connection failed: {error_message}")
|
||
return False, error_message
|
||
except Exception as e:
|
||
error_message = f"Unexpected error while checking MySQL connection: {str(e)}"
|
||
logger.error(f"✗ MySQL connection check failed: {error_message}")
|
||
return False, error_message
|
||
finally:
|
||
if temp_connection:
|
||
temp_connection.close()
|
||
|
||
@staticmethod
|
||
def check_database_exists(db_name: str, connection=None) -> bool:
|
||
"""
|
||
Check if a database exists
|
||
|
||
Args:
|
||
db_name: Database name to check
|
||
connection: Optional existing MySQL connection
|
||
|
||
Returns:
|
||
True if database exists, False otherwise
|
||
"""
|
||
temp_connection = None
|
||
try:
|
||
if connection is None:
|
||
# Create temporary connection without specifying database
|
||
temp_connection = pymysql.connect(
|
||
host=settings.MYSQL_HOST,
|
||
port=settings.MYSQL_PORT,
|
||
user=settings.MYSQL_USER,
|
||
password=settings.MYSQL_PASSWORD,
|
||
charset=settings.MYSQL_CHARSET,
|
||
cursorclass=pymysql.cursors.DictCursor
|
||
)
|
||
conn = temp_connection
|
||
else:
|
||
conn = connection
|
||
|
||
with conn.cursor() as cursor:
|
||
cursor.execute("SHOW DATABASES LIKE %s", (db_name,))
|
||
result = cursor.fetchone()
|
||
return result is not None
|
||
except Exception as e:
|
||
logger.error(f"Error checking database existence for {db_name}: {e}")
|
||
return False
|
||
finally:
|
||
if temp_connection:
|
||
temp_connection.close()
|
||
|
||
@staticmethod
|
||
def check_databases_exist(db_configs: List[DatabaseConfig]) -> Tuple[bool, List[str]]:
|
||
"""
|
||
Check if all configured databases exist
|
||
Each database config can use its own MySQL connection info
|
||
|
||
Args:
|
||
db_configs: List of DatabaseConfig objects
|
||
|
||
Returns:
|
||
Tuple of (all_exist: bool, missing_databases: List[str])
|
||
"""
|
||
if not db_configs:
|
||
return True, []
|
||
|
||
missing_databases = []
|
||
|
||
# Group database configs by connection info (host:port:user)
|
||
# This allows us to check multiple databases on the same server efficiently
|
||
connection_groups: Dict[Tuple[str, int, str], List[DatabaseConfig]] = {}
|
||
for db_config in db_configs:
|
||
# Use database-specific connection info if provided, otherwise use .env defaults
|
||
host = db_config.mysql_host if db_config.mysql_host is not None else settings.MYSQL_HOST
|
||
port = db_config.mysql_port if db_config.mysql_port is not None else settings.MYSQL_PORT
|
||
user = db_config.mysql_user if db_config.mysql_user is not None else settings.MYSQL_USER
|
||
connection_key = (host, port, user)
|
||
|
||
if connection_key not in connection_groups:
|
||
connection_groups[connection_key] = []
|
||
connection_groups[connection_key].append(db_config)
|
||
|
||
# Check databases for each connection group
|
||
for (host, port, user), configs in connection_groups.items():
|
||
temp_connection = None
|
||
try:
|
||
# Get password and charset from first config in group (or use defaults)
|
||
# Note: All configs in a group should use same password/charset for same host:port:user
|
||
first_config = configs[0]
|
||
password = first_config.mysql_password if first_config.mysql_password is not None else settings.MYSQL_PASSWORD
|
||
charset = first_config.mysql_charset if first_config.mysql_charset is not None else settings.MYSQL_CHARSET
|
||
|
||
logger.info(f"Checking databases on MySQL server {host}:{port} (user: {user})...")
|
||
temp_connection = pymysql.connect(
|
||
host=host,
|
||
port=port,
|
||
user=user,
|
||
password=password,
|
||
charset=charset,
|
||
cursorclass=pymysql.cursors.DictCursor,
|
||
connect_timeout=10
|
||
)
|
||
|
||
# Get all existing databases on this server
|
||
with temp_connection.cursor() as cursor:
|
||
cursor.execute("SHOW DATABASES")
|
||
existing_databases = {row['Database'] for row in cursor.fetchall()}
|
||
|
||
# Check each configured database in this group
|
||
for db_config in configs:
|
||
if db_config.database not in existing_databases:
|
||
missing_databases.append(
|
||
f"{db_config.database} (name: {db_config.name}, server: {host}:{port})"
|
||
)
|
||
else:
|
||
logger.info(f"✓ Database {db_config.database} exists on {host}:{port}")
|
||
|
||
except Exception as e:
|
||
# If we can't connect to check, mark all databases in this group as missing
|
||
logger.error(f"Error checking databases on {host}:{port}: {e}")
|
||
for db_config in configs:
|
||
missing_databases.append(
|
||
f"{db_config.database} (name: {db_config.name}, server: {host}:{port}, error: {str(e)})"
|
||
)
|
||
finally:
|
||
if temp_connection:
|
||
temp_connection.close()
|
||
|
||
return len(missing_databases) == 0, missing_databases
|
||
|
||
def _switch_database(self):
|
||
"""Switch to the target database using existing connection"""
|
||
try:
|
||
with self.connection.cursor() as cursor:
|
||
cursor.execute(f"USE `{self.db_config.database}`")
|
||
logger.info(f"Switched to MySQL database: {self.db_config.database} (name: {self.db_config.name})")
|
||
except Exception as e:
|
||
logger.error(f"Failed to switch to MySQL database {self.db_config.name}: {e}")
|
||
raise
|
||
|
||
def close(self):
|
||
"""Close MySQL connection"""
|
||
if self.connection:
|
||
self.connection.close()
|
||
logger.info("MySQL connection closed")
|
||
|
||
def _ensure_connection(self):
|
||
"""Ensure MySQL connection is alive and reconnect if necessary"""
|
||
max_retries = 3
|
||
for attempt in range(max_retries):
|
||
try:
|
||
# Check if connection is None
|
||
if self.connection is None:
|
||
logger.warning(f"MySQL connection for {self.db_config.name} is None, reconnecting... (attempt {attempt + 1}/{max_retries})")
|
||
self._connect()
|
||
# Verify connection was created
|
||
if self.connection is None:
|
||
raise RuntimeError(f"Failed to create connection for {self.db_config.name}")
|
||
return
|
||
|
||
# Check if connection is open
|
||
if not hasattr(self.connection, 'open') or not self.connection.open:
|
||
logger.warning(f"MySQL connection for {self.db_config.name} is closed, reconnecting... (attempt {attempt + 1}/{max_retries})")
|
||
try:
|
||
self.connection.close()
|
||
except Exception:
|
||
pass
|
||
self.connection = None
|
||
self._connect()
|
||
# Verify connection was created
|
||
if self.connection is None:
|
||
raise RuntimeError(f"Failed to create connection for {self.db_config.name}")
|
||
return
|
||
|
||
# Test connection with ping
|
||
# Note: ping(reconnect=True) may not work if connection is in bad state
|
||
try:
|
||
self.connection.ping(reconnect=False)
|
||
except Exception:
|
||
# Ping failed, connection is bad, recreate it
|
||
logger.warning(f"Ping failed for {self.db_config.name}, recreating connection...")
|
||
try:
|
||
self.connection.close()
|
||
except Exception:
|
||
pass
|
||
self.connection = None
|
||
self._connect()
|
||
if self.connection is None:
|
||
raise RuntimeError(f"Failed to create connection for {self.db_config.name}")
|
||
return
|
||
|
||
return # Connection is good
|
||
except (pymysql.err.OperationalError, pymysql.err.InterfaceError, AttributeError) as e:
|
||
# Connection is dead or invalid, try to reconnect
|
||
logger.warning(f"MySQL connection for {self.db_config.name} is dead/invalid, reconnecting... (attempt {attempt + 1}/{max_retries}): {e}")
|
||
try:
|
||
if self.connection:
|
||
self.connection.close()
|
||
except Exception:
|
||
pass
|
||
self.connection = None
|
||
if attempt < max_retries - 1:
|
||
import time
|
||
time.sleep(0.5 * (attempt + 1)) # Exponential backoff
|
||
else:
|
||
# Last attempt, raise the error
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Unexpected error checking MySQL connection for {self.db_config.name}: {e}")
|
||
# Set connection to None so next attempt will recreate it
|
||
try:
|
||
if self.connection:
|
||
self.connection.close()
|
||
except Exception:
|
||
pass
|
||
self.connection = None
|
||
if attempt < max_retries - 1:
|
||
import time
|
||
time.sleep(0.5 * (attempt + 1))
|
||
else:
|
||
raise
|
||
|
||
def fetch_all_documents(self) -> List[Dict]:
|
||
"""
|
||
Fetch all documents from MySQL table
|
||
|
||
Returns:
|
||
List of document dictionaries with database source info
|
||
"""
|
||
try:
|
||
# Ensure connection is alive before using it
|
||
self._ensure_connection()
|
||
|
||
# Execute query with retry mechanism for connection issues
|
||
max_retries = 3
|
||
results = None
|
||
for attempt in range(max_retries):
|
||
try:
|
||
# Ensure connection is still alive (may have been lost during previous attempt)
|
||
if attempt > 0:
|
||
self._ensure_connection()
|
||
else:
|
||
# Also ensure connection on first attempt (may have been set to None elsewhere)
|
||
self._ensure_connection()
|
||
|
||
# Double-check connection is not None before using
|
||
if self.connection is None:
|
||
logger.warning(f"Connection is None for {self.db_config.name} after _ensure_connection, retrying...")
|
||
if attempt < max_retries - 1:
|
||
import time
|
||
time.sleep(0.5 * (attempt + 1))
|
||
continue
|
||
else:
|
||
raise RuntimeError(f"Failed to establish connection for {self.db_config.name} after {max_retries} attempts")
|
||
|
||
# Ensure we're using the correct database
|
||
with self.connection.cursor() as cursor:
|
||
cursor.execute(f"USE `{self.db_config.database}`")
|
||
|
||
with self.connection.cursor() as cursor:
|
||
# Get actual columns in the table
|
||
cursor.execute(f"DESCRIBE `{self.db_config.table_name}`")
|
||
table_columns = {row['Field'] for row in cursor.fetchall()}
|
||
|
||
# Build query based on configured columns, only include columns that exist
|
||
columns = []
|
||
# Add ID column (required)
|
||
if self.db_config.id_column in table_columns:
|
||
columns.append(self.db_config.id_column)
|
||
else:
|
||
raise ValueError(f"ID column '{self.db_config.id_column}' does not exist in table {self.db_config.table_name}")
|
||
|
||
# 添加所有 content 列（支持多个列），只添加存在的列
|
||
for col in self.db_config.content_columns:
|
||
if col in table_columns:
|
||
columns.append(col)
|
||
else:
|
||
logger.warning(f"Content column '{col}' does not exist in table {self.db_config.table_name}, skipping")
|
||
|
||
if not columns or all(col == self.db_config.id_column for col in columns):
|
||
raise ValueError(f"No valid content columns found in table {self.db_config.table_name}")
|
||
|
||
# Add title column if configured and exists
|
||
if self.db_config.title_column and self.db_config.title_column in table_columns:
|
||
columns.append(self.db_config.title_column)
|
||
elif self.db_config.title_column:
|
||
logger.warning(f"Title column '{self.db_config.title_column}' does not exist in table {self.db_config.table_name}, skipping")
|
||
|
||
# Add metadata columns if specified, only add existing ones
|
||
if self.db_config.metadata_columns:
|
||
metadata_cols = [col.strip() for col in self.db_config.metadata_columns.split(",")]
|
||
for col in metadata_cols:
|
||
if col in table_columns:
|
||
columns.append(col)
|
||
else:
|
||
logger.warning(f"Metadata column '{col}' does not exist in table {self.db_config.table_name}, skipping")
|
||
|
||
query = f"SELECT {', '.join(columns)} FROM `{self.db_config.table_name}`"
|
||
cursor.execute(query)
|
||
results = cursor.fetchall()
|
||
break # Success, exit retry loop
|
||
|
||
except (pymysql.err.OperationalError, pymysql.err.InterfaceError) as e:
|
||
# Safely get error code
|
||
error_code = None
|
||
error_msg = str(e)
|
||
if hasattr(e, 'args') and e.args:
|
||
try:
|
||
error_code = e.args[0] if len(e.args) > 0 else None
|
||
error_msg = e.args[1] if len(e.args) > 1 else str(e)
|
||
except (IndexError, TypeError):
|
||
error_code = None
|
||
|
||
# Check for connection-related errors that should trigger reconnection
|
||
should_reconnect = False
|
||
if error_code == 2013: # Lost connection to MySQL server during query
|
||
should_reconnect = True
|
||
elif "Packet sequence number wrong" in error_msg or "Lost connection" in error_msg:
|
||
should_reconnect = True
|
||
|
||
if should_reconnect:
|
||
logger.warning(f"Connection error during query for {self.db_config.name} ({error_msg}), reconnecting and retrying... (attempt {attempt + 1}/{max_retries})")
|
||
# Close and recreate connection
|
||
try:
|
||
if self.connection:
|
||
self.connection.close()
|
||
except Exception:
|
||
pass
|
||
self.connection = None
|
||
|
||
if attempt < max_retries - 1:
|
||
import time
|
||
time.sleep(0.5 * (attempt + 1)) # Exponential backoff
|
||
continue
|
||
else:
|
||
logger.error(f"Failed to fetch documents after {max_retries} attempts: {e}")
|
||
raise
|
||
else:
|
||
# Other operational error, raise immediately
|
||
raise
|
||
except Exception as e:
|
||
# Non-connection errors, raise immediately
|
||
raise
|
||
|
||
if results is None:
|
||
raise RuntimeError(f"Failed to fetch documents from {self.db_config.name} after {max_retries} attempts")
|
||
logger.info(f"Fetched {len(results)} documents from MySQL database: {self.db_config.database} (name: {self.db_config.name})")
|
||
|
||
# 如果配置了文件字段，则需根据文件字段的值来获取文件内容
|
||
if self.db_config.file_column and self.db_config.file_column in self.db_config.content_columns:
|
||
# 文件处理成功的文档
|
||
processed_results = []
|
||
# 失败文档ID列表
|
||
failed_doc_ids = set()
|
||
logger.info(
|
||
f"MySQL database: {self.db_config.database} (table: {self.db_config.table_name}) has set file_column(value: {self.db_config.file_column}) field.")
|
||
for doc in results:
|
||
# 获取每个doc的文件字段值，即文件标识符
|
||
file_column_value = doc.get(self.db_config.file_column)
|
||
if not file_column_value or file_column_value is None:
|
||
processed_results.append(doc)
|
||
continue
|
||
file_identifiers = [file_identifier.strip() for file_identifier in file_column_value.split(",")]
|
||
try:
|
||
file_content_parts = []
|
||
for file_identifier in file_identifiers:
|
||
# 调用接口获取文件字节流及文件名
|
||
file_byte, file_name = self.fetch_file_bytes(file_identifier)
|
||
if not file_byte:
|
||
logger.warning(
|
||
f"文档ID:{doc.get(self.db_config.id_column, '未知')}，文件标识符：{file_identifier} 获取字节流失败")
|
||
failed_doc_ids.add(doc.get(self.db_config.id_column))
|
||
break
|
||
# 调用文件读取方法获取文本内容
|
||
file_text = self.file_parser.get_text_from_bytes(file_byte, file_name)
|
||
if file_text and file_text.strip():
|
||
file_content_parts.append(file_text.strip())
|
||
logger.info(f"成功读取文件内容: {file_identifier} (长度: {len(file_text)})")
|
||
else:
|
||
logger.warning(f"文件标识符 {file_identifier} 提取的文本内容为空")
|
||
failed_doc_ids.add(doc.get(self.db_config.id_column))
|
||
break
|
||
if file_content_parts:
|
||
file_content = "\n".join(file_content_parts)
|
||
# 将文件内容添加到文档中
|
||
doc[self.db_config.file_column] = file_content
|
||
logger.info(
|
||
f"文档 {doc.get(self.db_config.id_column, '未知')} 成功合并 {len(file_content_parts)} 个文件内容")
|
||
# 将处理成功的文档添加到新列表中
|
||
processed_results.append(doc)
|
||
else:
|
||
logger.warning(
|
||
f"文档 {doc.get(self.db_config.id_column, '未知')} 没有成功提取任何文件内容")
|
||
failed_doc_ids.add(doc.get(self.db_config.id_column))
|
||
|
||
except Exception as e:
|
||
logger.error(
|
||
f"处理文档({doc.get(self.db_config.id_column, '未知')}) 文件字段: {doc[self.db_config.file_column]} 时出错: {e}")
|
||
failed_doc_ids.add(doc.get(self.db_config.id_column))
|
||
continue # 继续下一次doc循环
|
||
|
||
# 记录失败文档统计信息
|
||
if failed_doc_ids:
|
||
failed_doc_ids_str = {str(doc_id) for doc_id in failed_doc_ids}
|
||
ids_str = ', '.join(sorted(failed_doc_ids_str))
|
||
logger.info(f"数据库配置[{self.db_config.name}]本次同步在处理文件过程中共有 {len(failed_doc_ids)} 个文档因文件处理失败被跳过: [{ids_str}]")
|
||
else:
|
||
logger.info(f"数据库配置[{self.db_config.name}]所有文档文件处理成功")
|
||
|
||
results = processed_results
|
||
|
||
# Add database source information to each document
|
||
if results:
|
||
for doc in results:
|
||
doc['_db_source'] = self.db_config.name
|
||
doc['_db_database'] = self.db_config.database
|
||
doc['_db_table'] = self.db_config.table_name
|
||
|
||
return results
|
||
except Exception as e:
|
||
logger.error(f"Error fetching documents from MySQL database {self.db_config.name}: {e}")
|
||
raise
|
||
|
||
def fetch_new_documents(self, last_sync_time: Optional[datetime] = None) -> List[Dict]:
|
||
"""
|
||
Fetch new or updated documents since last sync
|
||
|
||
Args:
|
||
last_sync_time: Last synchronization timestamp
|
||
|
||
Returns:
|
||
List of new/updated document dictionaries with database source info
|
||
"""
|
||
try:
|
||
# Ensure connection is alive before using it
|
||
self._ensure_connection()
|
||
|
||
# Ensure we're using the correct database
|
||
with self.connection.cursor() as cursor:
|
||
cursor.execute(f"USE `{self.db_config.database}`")
|
||
|
||
with self.connection.cursor() as cursor:
|
||
# Get actual columns in the table
|
||
cursor.execute(f"DESCRIBE `{self.db_config.table_name}`")
|
||
table_columns = {row['Field'] for row in cursor.fetchall()}
|
||
|
||
# Build query based on configured columns, only include columns that exist
|
||
columns = []
|
||
# Add ID column (required)
|
||
if self.db_config.id_column in table_columns:
|
||
columns.append(self.db_config.id_column)
|
||
else:
|
||
raise ValueError(f"ID column '{self.db_config.id_column}' does not exist in table {self.db_config.table_name}")
|
||
|
||
# 添加所有 content 列（支持多个列），只添加存在的列
|
||
for col in self.db_config.content_columns:
|
||
if col in table_columns:
|
||
columns.append(col)
|
||
else:
|
||
logger.warning(f"Content column '{col}' does not exist in table {self.db_config.table_name}, skipping")
|
||
|
||
if not columns or all(col == self.db_config.id_column for col in columns):
|
||
raise ValueError(f"No valid content columns found in table {self.db_config.table_name}")
|
||
|
||
# Add title column if configured and exists
|
||
if self.db_config.title_column and self.db_config.title_column in table_columns:
|
||
columns.append(self.db_config.title_column)
|
||
elif self.db_config.title_column:
|
||
logger.warning(f"Title column '{self.db_config.title_column}' does not exist in table {self.db_config.table_name}, skipping")
|
||
|
||
# Add metadata columns if specified, only add existing ones
|
||
if self.db_config.metadata_columns:
|
||
metadata_cols = [col.strip() for col in self.db_config.metadata_columns.split(",")]
|
||
for col in metadata_cols:
|
||
if col in table_columns:
|
||
columns.append(col)
|
||
else:
|
||
logger.warning(f"Metadata column '{col}' does not exist in table {self.db_config.table_name}, skipping")
|
||
|
||
query = f"SELECT {', '.join(columns)} FROM `{self.db_config.table_name}`"
|
||
|
||
# Use updated_at_column if configured and last_sync_time is provided
|
||
# Otherwise, fetch all documents (no incremental sync)
|
||
if last_sync_time and self.db_config.updated_at_column:
|
||
# Use configured updated_at column for incremental sync
|
||
# MySQL DATETIME only has second precision, so we need to ensure
|
||
# the comparison works correctly. Use >= to include records with same timestamp
|
||
# and format the datetime to match MySQL's DATETIME format
|
||
query += f" WHERE {self.db_config.updated_at_column} > %s"
|
||
# Format datetime to MySQL DATETIME format (YYYY-MM-DD HH:MM:SS)
|
||
# This ensures proper comparison with MySQL DATETIME type
|
||
mysql_datetime = last_sync_time.strftime('%Y-%m-%d %H:%M:%S')
|
||
cursor.execute(query, (mysql_datetime,))
|
||
logger.info(
|
||
f"Incremental sync using {self.db_config.updated_at_column}: "
|
||
f"last_sync_time={last_sync_time} (formatted: {mysql_datetime}), "
|
||
f"query: {query}"
|
||
)
|
||
else:
|
||
# No updated_at column configured or no last_sync_time, fetch all
|
||
if last_sync_time and not self.db_config.updated_at_column:
|
||
logger.warning(
|
||
f"No updated_at_column configured for {self.db_config.name}, "
|
||
"fetching all documents for incremental sync"
|
||
)
|
||
cursor.execute(query)
|
||
|
||
results = cursor.fetchall()
|
||
|
||
# 如果配置了文件字段，则需根据文件字段的值来获取文件内容
|
||
if self.db_config.file_column and self.db_config.file_column in self.db_config.content_columns:
|
||
# 文件处理成功的文档
|
||
processed_results = []
|
||
# 失败文档ID列表
|
||
failed_doc_ids = set()
|
||
logger.info(
|
||
f"MySQL database: {self.db_config.database} (table: {self.db_config.table_name}) has set file_column(value: {self.db_config.file_column}) field.")
|
||
for doc in results:
|
||
# 获取每个doc的文件字段值，即文件标识符
|
||
file_column_value = doc.get(self.db_config.file_column)
|
||
if not file_column_value or file_column_value is None:
|
||
processed_results.append(doc)
|
||
continue
|
||
file_identifiers = [file_identifier.strip() for file_identifier in file_column_value.split(",")]
|
||
try:
|
||
file_content_parts = []
|
||
for file_identifier in file_identifiers:
|
||
# 调用接口获取文件字节流及文件名
|
||
file_byte, file_name = self.fetch_file_bytes(file_identifier)
|
||
if not file_byte:
|
||
logger.warning(
|
||
f"文档ID:{doc.get(self.db_config.id_column, '未知')}，文件标识符：{file_identifier} 获取字节流失败")
|
||
failed_doc_ids.add(doc.get(self.db_config.id_column))
|
||
break
|
||
# 调用文件读取方法获取文本内容
|
||
file_text = self.file_parser.get_text_from_bytes(file_byte, file_name)
|
||
if file_text and file_text.strip():
|
||
file_content_parts.append(file_text.strip())
|
||
logger.info(f"成功读取文件内容: {file_identifier} (长度: {len(file_text)})")
|
||
else:
|
||
logger.warning(f"文件标识符 {file_identifier} 提取的文本内容为空")
|
||
failed_doc_ids.add(doc.get(self.db_config.id_column))
|
||
break
|
||
if file_content_parts:
|
||
file_content = "\n".join(file_content_parts)
|
||
# 将文件内容添加到文档中
|
||
doc[self.db_config.file_column] = file_content
|
||
logger.info(
|
||
f"文档 {doc.get(self.db_config.id_column, '未知')} 成功合并 {len(file_content_parts)} 个文件内容")
|
||
# 将处理成功的文档添加到新列表中
|
||
processed_results.append(doc)
|
||
else:
|
||
logger.warning(
|
||
f"文档 {doc.get(self.db_config.id_column, '未知')} 没有成功提取任何文件内容")
|
||
failed_doc_ids.add(doc.get(self.db_config.id_column))
|
||
|
||
except Exception as e:
|
||
logger.error(
|
||
f"处理文档({doc.get(self.db_config.id_column, '未知')}) 文件字段: {doc[self.db_config.file_column]} 时出错: {e}")
|
||
failed_doc_ids.add(doc.get(self.db_config.id_column))
|
||
continue # 继续下一次doc循环
|
||
|
||
# 记录失败文档统计信息
|
||
if failed_doc_ids:
|
||
failed_doc_ids_str = {str(doc_id) for doc_id in failed_doc_ids}
|
||
ids_str = ', '.join(sorted(failed_doc_ids_str))
|
||
logger.warning(
|
||
f"数据库配置[{self.db_config.name}]本次同步在处理文件过程中共有 {len(failed_doc_ids)} 个文档因文件处理失败被跳过: [{ids_str}]")
|
||
else:
|
||
logger.info(f"数据库配置[{self.db_config.name}]所有文档文件处理成功")
|
||
|
||
results = processed_results
|
||
|
||
# Add database source information to each document
|
||
if results:
|
||
for doc in results:
|
||
doc['_db_source'] = self.db_config.name
|
||
doc['_db_database'] = self.db_config.database
|
||
doc['_db_table'] = self.db_config.table_name
|
||
|
||
logger.info(f"Fetched {len(results)} new/updated documents from MySQL database: {self.db_config.database} (name: {self.db_config.name})")
|
||
return results
|
||
except Exception as e:
|
||
logger.error(f"Error fetching new documents from MySQL database {self.db_config.name}: {e}")
|
||
raise
|
||
|
||
def get_document_by_id(self, doc_id: str) -> Optional[Dict]:
|
||
"""
|
||
Get a single document by ID
|
||
|
||
Args:
|
||
doc_id: Document ID
|
||
|
||
Returns:
|
||
Document dictionary or None
|
||
"""
|
||
try:
|
||
# Ensure we're using the correct database
|
||
with self.connection.cursor() as cursor:
|
||
cursor.execute(f"USE `{self.db_config.database}`")
|
||
|
||
with self.connection.cursor() as cursor:
|
||
# Get actual columns in the table
|
||
cursor.execute(f"DESCRIBE `{self.db_config.table_name}`")
|
||
table_columns = {row['Field'] for row in cursor.fetchall()}
|
||
|
||
# Build query based on configured columns, only include columns that exist
|
||
columns = []
|
||
# Add ID column (required)
|
||
if self.db_config.id_column in table_columns:
|
||
columns.append(self.db_config.id_column)
|
||
else:
|
||
raise ValueError(f"ID column '{self.db_config.id_column}' does not exist in table {self.db_config.table_name}")
|
||
|
||
# 添加所有 content 列（支持多个列），只添加存在的列
|
||
for col in self.db_config.content_columns:
|
||
if col in table_columns:
|
||
columns.append(col)
|
||
else:
|
||
logger.warning(f"Content column '{col}' does not exist in table {self.db_config.table_name}, skipping")
|
||
|
||
if not columns or all(col == self.db_config.id_column for col in columns):
|
||
raise ValueError(f"No valid content columns found in table {self.db_config.table_name}")
|
||
|
||
# Add title column if configured and exists
|
||
if self.db_config.title_column and self.db_config.title_column in table_columns:
|
||
columns.append(self.db_config.title_column)
|
||
elif self.db_config.title_column:
|
||
logger.warning(f"Title column '{self.db_config.title_column}' does not exist in table {self.db_config.table_name}, skipping")
|
||
|
||
# Add metadata columns if specified, only add existing ones
|
||
if self.db_config.metadata_columns:
|
||
metadata_cols = [col.strip() for col in self.db_config.metadata_columns.split(",")]
|
||
for col in metadata_cols:
|
||
if col in table_columns:
|
||
columns.append(col)
|
||
else:
|
||
logger.warning(f"Metadata column '{col}' does not exist in table {self.db_config.table_name}, skipping")
|
||
|
||
query = f"SELECT {', '.join(columns)} FROM `{self.db_config.table_name}` WHERE {self.db_config.id_column} = %s"
|
||
cursor.execute(query, (doc_id,))
|
||
result = cursor.fetchone()
|
||
if result:
|
||
result['_db_source'] = self.db_config.name
|
||
result['_db_database'] = self.db_config.database
|
||
result['_db_table'] = self.db_config.table_name
|
||
return result
|
||
except Exception as e:
|
||
logger.error(f"Error fetching document {doc_id} from MySQL database {self.db_config.name}: {e}")
|
||
return None
|
||
|
||
# 从API获取文件字节流及文件名
|
||
def fetch_file_bytes(self, file_identifier: str) -> tuple[bytes, str]:
|
||
"""
|
||
从API获取文件流、文件名
|
||
Args:
|
||
file_identifier: 文件标识值 J6pLOFxe
|
||
Returns:
|
||
tuple: (文件字节流, 文件名)
|
||
"""
|
||
# base_url = "http://172.20.32.184:8000/api/file/open/downloadByIdentifier"
|
||
base_url = settings.FILE_DOWNLOAD_BASE_URL
|
||
url = f"{base_url}/{file_identifier}"
|
||
try:
|
||
# 发送GET请求获取文件流
|
||
response = requests.get(url, stream=True, timeout=30)
|
||
|
||
# 检查响应内容是否为文件
|
||
content_type = response.headers.get('Content-Type', '')
|
||
if 'application/json' in content_type:
|
||
# 如果是JSON响应，尝试解析错误信息
|
||
error_data = response.json()
|
||
error_msg = error_data.get('msg', '未知错误')
|
||
error_code = error_data.get('code', '未知')
|
||
raise ValueError(f"API返回错误: {error_msg} (代码: {error_code})")
|
||
response.raise_for_status() # 如果请求失败，抛出HTTPError异常
|
||
# 获取文件字节流
|
||
file_bytes = response.content
|
||
logger.info(f"Successfully fetch file bytes, file: {file_identifier} (size: {len(file_bytes)} bytes)")
|
||
|
||
# 从响应头中获取文件名
|
||
file_name = "unknown"
|
||
|
||
# 从 Content-Disposition 头获取文件名
|
||
content_disposition = response.headers.get('Content-Disposition', '')
|
||
if 'filename=' in content_disposition:
|
||
# 提取文件名，处理可能的引号和多个filename参数
|
||
import re
|
||
# match = re.search(r'filename=["\']?([^"\']+)["\']?', content_disposition)
|
||
match = re.search(r'filename=["\']?([^;"\']+)["\']?', content_disposition)
|
||
if match:
|
||
encoded_filename = match.group(1)
|
||
# 解码URL编码的中文文件名
|
||
file_name = urllib.parse.unquote(encoded_filename)
|
||
|
||
return file_bytes, file_name
|
||
|
||
except requests.exceptions.RequestException as e:
|
||
logger.error(f"Failed to fetch file bytes, file: {file_identifier}: {e}")
|
||
raise |