284 lines
10 KiB
Python
284 lines
10 KiB
Python
"""MySQL synchronization implementation"""
|
|
import urllib
|
|
import pymysql
|
|
import requests
|
|
from typing import List, Dict, Optional, Tuple, Any, Set
|
|
from datetime import datetime
|
|
from loguru import logger
|
|
from config import DatabaseDataSourceConfig as DatabaseConfig, settings
|
|
from rag.file_parser import FileParser
|
|
from sync.base_sync import BaseSync
|
|
|
|
|
|
class MySQLSync(BaseSync):
|
|
"""Handle synchronization between MySQL and ChromaDB"""
|
|
|
|
def __init__(self, db_config: DatabaseConfig, vector_store_manager=None):
|
|
"""
|
|
Initialize MySQL sync with database configuration
|
|
|
|
Args:
|
|
db_config: DatabaseConfig object containing database table info
|
|
"""
|
|
super().__init__(db_config)
|
|
self.db_config = self.config
|
|
self.vector_store_manager = vector_store_manager
|
|
self.file_parser = FileParser()
|
|
self.connection = None
|
|
|
|
def _connect(self):
|
|
"""Create MySQL connection"""
|
|
# Get connection parameters from config or use defaults
|
|
host = self.db_config.host or settings.HOST
|
|
port = self.db_config.port or settings.PORT
|
|
user = self.db_config.user or settings.USER
|
|
password = self.db_config.password or settings.PASSWORD
|
|
|
|
# Log connection info (without password)
|
|
logger.info(f"Connecting to MySQL database: {self.db_config.database}")
|
|
logger.debug(f"MySQL connection details: host={host}, port={port}, user={user}")
|
|
|
|
# Create connection
|
|
self.connection = pymysql.connect(
|
|
host=host,
|
|
port=port,
|
|
user=user,
|
|
password=password,
|
|
database=self.db_config.database
|
|
)
|
|
|
|
def fetch_all_documents(self, last_sync_time=None) -> List[Dict[str, Any]]:
|
|
"""
|
|
Fetch documents from the MySQL database
|
|
|
|
Args:
|
|
last_sync_time: Last synchronization time (for incremental sync)
|
|
|
|
Returns:
|
|
List of documents
|
|
"""
|
|
self._connect()
|
|
cursor = self.connection.cursor(pymysql.cursors.DictCursor)
|
|
try:
|
|
# Build query with all columns
|
|
columns = [
|
|
self.db_config.id_column,
|
|
self.db_config.title_column,
|
|
self.db_config.content_column,
|
|
self.db_config.file_column,
|
|
self.db_config.updated_at_column
|
|
] if self.db_config.updated_at_column else [
|
|
self.db_config.id_column,
|
|
self.db_config.title_column,
|
|
self.db_config.content_column,
|
|
self.db_config.file_column
|
|
]
|
|
|
|
# Remove duplicates and None values
|
|
columns = list(set([col for col in columns if col]))
|
|
|
|
# Build the query
|
|
query = f"SELECT {', '.join(columns)} FROM {self.db_config.table_name}"
|
|
|
|
# Add incremental sync condition if applicable
|
|
params = []
|
|
if last_sync_time and self.db_config.updated_at_column:
|
|
query += f" WHERE {self.db_config.updated_at_column} > %s"
|
|
params.append(last_sync_time)
|
|
|
|
logger.debug(f"MySQL query: {query}, params: {params}")
|
|
cursor.execute(query, params)
|
|
|
|
# Parse results
|
|
documents = []
|
|
for row in cursor.fetchall():
|
|
# Handle file content if file_column is specified
|
|
if self.db_config.file_column and row.get(self.db_config.file_column):
|
|
# Extract file path from the file column
|
|
file_path = row[self.db_config.file_column]
|
|
|
|
# Load file content if file source is configured
|
|
if self.db_config.file_source_type:
|
|
file_content = self._load_file_content(file_path)
|
|
if file_content:
|
|
row[self.db_config.content_column] = file_content
|
|
|
|
# Generate unique document ID
|
|
record_id = str(row[self.db_config.id_column])
|
|
doc_id = self.generate_doc_id(record_id)
|
|
row['id'] = doc_id
|
|
|
|
# Check if document has already been synced
|
|
if self.vector_store_manager and self.vector_store_manager.document_exists(doc_id + "_chunk_0"):
|
|
# If document is already synced, check if it's been modified since last sync
|
|
if last_sync_time and self.db_config.updated_at_column and row.get(self.db_config.updated_at_column):
|
|
# Skip if not modified since last sync
|
|
if row[self.db_config.updated_at_column] <= last_sync_time:
|
|
continue
|
|
elif last_sync_time:
|
|
# No updated_at column, skip since we can't determine if it's been modified
|
|
continue
|
|
|
|
documents.append(row)
|
|
|
|
return documents
|
|
finally:
|
|
cursor.close()
|
|
|
|
def fetch_new_documents(self, last_sync_time=None) -> List[Dict[str, Any]]:
|
|
"""
|
|
Fetch new/updated documents from the MySQL database since last sync time
|
|
|
|
Args:
|
|
last_sync_time: Last synchronization time
|
|
|
|
Returns:
|
|
List of new/updated documents
|
|
"""
|
|
# Call fetch_all_documents which now handles document existence checks
|
|
return self.fetch_all_documents(last_sync_time)
|
|
|
|
def get_synced_document_ids(self) -> Set[str]:
|
|
"""
|
|
Get IDs of all documents in the MySQL database
|
|
|
|
Returns:
|
|
Set of document IDs
|
|
"""
|
|
cursor = self.connection.cursor(pymysql.cursors.DictCursor)
|
|
try:
|
|
query = f"SELECT {self.db_config.id_column} FROM {self.db_config.table_name}"
|
|
cursor.execute(query)
|
|
return {str(row[self.db_config.id_column]) for row in cursor.fetchall()}
|
|
finally:
|
|
cursor.close()
|
|
|
|
def generate_doc_id(self, record_id: str) -> str:
|
|
"""
|
|
Generate a unique document ID for MySQL records
|
|
|
|
Args:
|
|
record_id: ID of the record in the database
|
|
|
|
Returns:
|
|
Unique document ID based on database, table, and record ID
|
|
"""
|
|
# 为 MySQL 记录生成唯一的文档 ID
|
|
return f"{self.db_config.database}_{self.db_config.table_name}_{record_id}"
|
|
|
|
def doc_to_llamaindex_doc(self, doc: Dict) -> 'Document':
|
|
"""
|
|
Convert MySQL document to LlamaIndex Document
|
|
|
|
Args:
|
|
doc: MySQL document dictionary
|
|
|
|
Returns:
|
|
LlamaIndex Document object
|
|
"""
|
|
from llama_index.core import Document
|
|
|
|
# 处理多个 content 列（支持合并多个列的内容）
|
|
if self.db_config:
|
|
# 使用配置的多个 content 列
|
|
content_columns = self.db_config.content_columns
|
|
content_separator = self.db_config.content_separator
|
|
else:
|
|
# 向后兼容：使用单个 content_column
|
|
content_columns = ["content"] # Default to "content" column
|
|
content_separator = "\n"
|
|
|
|
# 合并所有 content 列的内容
|
|
content_parts = []
|
|
for col in content_columns:
|
|
col_value = doc.get(col, "")
|
|
if col_value:
|
|
content_parts.append(str(col_value))
|
|
|
|
# 用指定的分隔符连接多个列的内容
|
|
if content_parts:
|
|
if content_separator:
|
|
content = content_separator.join(content_parts)
|
|
else:
|
|
content = " ".join(content_parts) # 如果没有指定分隔符，使用空格
|
|
else:
|
|
content = ""
|
|
|
|
title = doc.get('title', "")
|
|
doc_id = doc.get('id', "")
|
|
|
|
# Build metadata
|
|
metadata = {
|
|
"doc_id": doc_id,
|
|
"source": "mysql",
|
|
"database": self.db_config.database,
|
|
"table": self.db_config.table_name
|
|
}
|
|
|
|
if title:
|
|
metadata["title"] = title
|
|
|
|
# Create Document
|
|
return Document(
|
|
text=content,
|
|
id_=doc_id,
|
|
metadata=metadata
|
|
)
|
|
|
|
|
|
|
|
@staticmethod
|
|
def check_data_source_exists(config: DatabaseConfig) -> bool:
|
|
"""
|
|
Check if the MySQL database exists and is accessible
|
|
|
|
Args:
|
|
config: Database configuration
|
|
|
|
Returns:
|
|
True if database exists and is accessible, False otherwise
|
|
"""
|
|
temp_connection = None
|
|
try:
|
|
# Get connection parameters
|
|
host = config.mysql_host or settings.MYSQL_HOST
|
|
port = config.mysql_port or settings.MYSQL_PORT
|
|
user = config.mysql_user or settings.MYSQL_USER
|
|
password = config.mysql_password or settings.MYSQL_PASSWORD
|
|
# Create connection
|
|
temp_connection = pymysql.connect(
|
|
host=host,
|
|
port=port,
|
|
user=user,
|
|
password=password
|
|
)
|
|
|
|
# Check if database exists
|
|
cursor = temp_connection.cursor()
|
|
cursor.execute(f"SHOW DATABASES LIKE '{config.database}'")
|
|
result = cursor.fetchone()
|
|
cursor.close()
|
|
|
|
if not result:
|
|
logger.error(f"Database {config.database} does not exist")
|
|
return False
|
|
|
|
# Check if table exists
|
|
temp_connection.select_db(config.database)
|
|
cursor = temp_connection.cursor()
|
|
cursor.execute(f"SHOW TABLES LIKE '{config.table_name}'")
|
|
result = cursor.fetchone()
|
|
cursor.close()
|
|
|
|
if not result:
|
|
logger.error(f"Table {config.table_name} does not exist in database {config.database}")
|
|
return False
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error checking MySQL data source: {e}")
|
|
return False
|
|
finally:
|
|
if temp_connection:
|
|
temp_connection.close()
|