476 lines
23 KiB
Python
476 lines
23 KiB
Python
"""
|
||
Background service for syncing MySQL data to ChromaDB
|
||
"""
|
||
import asyncio
|
||
from datetime import datetime
|
||
from typing import Set, List, Dict
|
||
from loguru import logger
|
||
from config import settings, DatabaseConfig
|
||
from database import MySQLSync
|
||
from rag import VectorStoreManager, DocumentProcessor
|
||
|
||
|
||
class SyncService:
|
||
"""Service for synchronizing MySQL data to ChromaDB"""
|
||
|
||
def __init__(self):
|
||
self.db_configs = settings.get_database_configs()
|
||
|
||
# Check if all configured databases exist before proceeding
|
||
# This will check each database using its own connection info (if specified)
|
||
logger.info("Checking if all configured databases exist...")
|
||
all_exist, missing_databases = MySQLSync.check_databases_exist(self.db_configs)
|
||
|
||
if not all_exist:
|
||
error_msg = (
|
||
f"Error: The following databases do not exist or cannot be accessed:\n"
|
||
f" {', '.join(missing_databases)}\n"
|
||
f"Please check:\n"
|
||
f" 1. All databases exist on their respective MySQL servers\n"
|
||
f" 2. MySQL servers are running and accessible\n"
|
||
f" 3. MySQL usernames and passwords are correct\n"
|
||
f" 4. Firewall rules allow connections to MySQL ports"
|
||
)
|
||
logger.error(error_msg)
|
||
raise RuntimeError(error_msg)
|
||
|
||
logger.info(f"✓ All {len(self.db_configs)} configured database(s) exist")
|
||
|
||
self.mysql_syncs: Dict[str, MySQLSync] = {}
|
||
self.vector_store_manager = VectorStoreManager()
|
||
self.document_processor = DocumentProcessor()
|
||
self.last_sync_times: Dict[str, datetime] = {}
|
||
self.synced_doc_ids: Dict[str, Set[str]] = {}
|
||
self._running = False
|
||
self._sync_in_progress = False # Flag to prevent concurrent syncs
|
||
self._auto_sync_task = None # Reference to auto sync task to prevent multiple instances
|
||
|
||
# Initialize MySQL connection - create separate connection for each database
|
||
# This avoids connection state issues when running in thread pool
|
||
for db_config in self.db_configs:
|
||
try:
|
||
# Create separate connection for each database to avoid connection state conflicts
|
||
# when running in thread pool (connections are not thread-safe)
|
||
mysql_sync = MySQLSync(db_config)
|
||
self.mysql_syncs[db_config.name] = mysql_sync
|
||
|
||
self.last_sync_times[db_config.name] = None
|
||
self.synced_doc_ids[db_config.name] = set()
|
||
except Exception as e:
|
||
logger.error(f"Failed to initialize MySQL sync for {db_config.name}: {e}")
|
||
raise
|
||
|
||
async def sync_all(self, force: bool = False):
|
||
"""
|
||
Sync all documents from all MySQL databases to ChromaDB
|
||
|
||
Args:
|
||
force: If True, re-process all documents even if they exist (default: False)
|
||
"""
|
||
# Prevent concurrent syncs
|
||
if self._sync_in_progress:
|
||
logger.warning("Sync already in progress, skipping this request")
|
||
return
|
||
|
||
self._sync_in_progress = True
|
||
sync_start_time = datetime.now()
|
||
try:
|
||
logger.info(f"[同步] 开始全量同步: {len(self.db_configs)} 个数据库")
|
||
|
||
# Run all synchronous operations in thread pool to avoid blocking event loop
|
||
import asyncio
|
||
loop = asyncio.get_event_loop()
|
||
|
||
def sync_work():
|
||
"""Synchronous work that runs in thread pool"""
|
||
# Get existing document IDs from ChromaDB to avoid re-processing
|
||
existing_doc_ids = set()
|
||
if not force:
|
||
existing_doc_ids = self.vector_store_manager.get_existing_doc_ids()
|
||
|
||
all_chunked_docs = [] # 存储所有分块后的文档
|
||
total_docs = 0
|
||
skipped_docs_count = 0
|
||
processed_databases = 0
|
||
|
||
# Sync from each database
|
||
for db_index, db_config in enumerate(self.db_configs, 1):
|
||
try:
|
||
mysql_sync = self.mysql_syncs[db_config.name]
|
||
mysql_docs = mysql_sync.fetch_all_documents()
|
||
|
||
if not mysql_docs:
|
||
processed_databases += 1
|
||
continue
|
||
|
||
# 从配置中获取内容列
|
||
content_columns = db_config.content_columns
|
||
|
||
# 获取每个数据源对应 ChromaDB metadata 中的 “content_column”
|
||
db_content_column_str = self.vector_store_manager.get_specific_db_source_metadata(db_config.database+'_'+db_config.table_name)
|
||
if db_content_column_str:
|
||
existing_db_content_columns = [col.strip() for col in db_content_column_str.split(",")]
|
||
# 获取 existing_doc_content_columns 与 content_columns 差异项
|
||
diff_columns = set(content_columns) - set(existing_db_content_columns)
|
||
else: # 说明 ChromaDB 中无该数据源的 metadata，需全量处理，无需过滤
|
||
diff_columns = set(content_columns)
|
||
|
||
# Filter out documents that already exist (if not forcing)
|
||
db_skipped_count = 0
|
||
if not force:
|
||
id_column = db_config.id_column
|
||
# 若 diff_columns为空，说明无新增列，则需过滤ChromaDB中已存在的doc
|
||
if not diff_columns:
|
||
new_mysql_docs = []
|
||
for doc in mysql_docs:
|
||
doc_id = str(doc.get(id_column, ""))
|
||
# 生成唯一文档标识符（数据库标识名称_表名_文档ID）{db_source}_{table_name}_{id}
|
||
unique_doc_id = f"{db_config.name}_{db_config.table_name}_{doc_id}"
|
||
if unique_doc_id not in existing_doc_ids:
|
||
new_mysql_docs.append(doc)
|
||
else:
|
||
db_skipped_count += 1
|
||
skipped_docs_count += 1
|
||
|
||
if not new_mysql_docs:
|
||
self.synced_doc_ids[db_config.name] = {
|
||
str(doc.get(id_column)) for doc in mysql_docs
|
||
}
|
||
self.last_sync_times[db_config.name] = datetime.now()
|
||
processed_databases += 1
|
||
continue
|
||
mysql_docs = new_mysql_docs
|
||
elif db_content_column_str:
|
||
self.vector_store_manager.delete_documents_by_source(db_config.database+'_'+db_config.table_name)
|
||
|
||
# Process and chunk documents
|
||
documents = self.document_processor.process_documents(mysql_docs, db_config)
|
||
chunked_docs = self.document_processor.chunk_documents(documents)
|
||
all_chunked_docs.extend(chunked_docs)
|
||
|
||
# Update synced document IDs for this database
|
||
id_column = db_config.id_column
|
||
self.synced_doc_ids[db_config.name] = {
|
||
str(doc.get(id_column)) for doc in mysql_docs
|
||
}
|
||
self.last_sync_times[db_config.name] = datetime.now()
|
||
total_docs += len(mysql_docs)
|
||
processed_databases += 1
|
||
|
||
except Exception as e:
|
||
logger.error(f"[同步] 数据库 {db_config.name} 同步失败: {e}")
|
||
processed_databases += 1
|
||
continue
|
||
|
||
return all_chunked_docs, total_docs, skipped_docs_count, processed_databases
|
||
|
||
# Run synchronous work in thread pool
|
||
all_chunked_docs, total_docs, skipped_docs_count, processed_databases = await loop.run_in_executor(None, sync_work)
|
||
|
||
# Add all documents to vector store (run in thread pool to avoid blocking event loop)
|
||
if all_chunked_docs:
|
||
await loop.run_in_executor(
|
||
None,
|
||
self.vector_store_manager.add_documents,
|
||
all_chunked_docs,
|
||
not force # skip_existing true
|
||
)
|
||
|
||
sync_duration = (datetime.now() - sync_start_time).total_seconds()
|
||
logger.info(
|
||
f"[同步] 完成: {processed_databases}/{len(self.db_configs)} 个数据库, "
|
||
f"{total_docs} 个新文档, {len(all_chunked_docs)} 个分块"
|
||
+ (f", 跳过 {skipped_docs_count} 个已存在" if skipped_docs_count > 0 else "")
|
||
+ f", 耗时 {sync_duration:.1f}秒"
|
||
)
|
||
else:
|
||
sync_duration = (datetime.now() - sync_start_time).total_seconds()
|
||
if skipped_docs_count > 0:
|
||
logger.info(f"[同步] 完成: 所有 {skipped_docs_count} 个文档已存在, 耗时 {sync_duration:.1f}秒")
|
||
else:
|
||
logger.warning(f"[同步] 没有文档需要同步, 耗时 {sync_duration:.1f}秒")
|
||
except Exception as e:
|
||
sync_duration = (datetime.now() - sync_start_time).total_seconds()
|
||
logger.error(f"[同步进度] ✗ 同步过程中出错 (耗时: {sync_duration:.1f} 秒): {e}")
|
||
raise
|
||
finally:
|
||
self._sync_in_progress = False
|
||
|
||
async def sync_incremental(self):
|
||
"""Sync only new/updated documents from MySQL to ChromaDB"""
|
||
# Prevent concurrent syncs
|
||
if self._sync_in_progress:
|
||
logger.warning("Sync already in progress, skipping incremental sync")
|
||
return
|
||
|
||
self._sync_in_progress = True
|
||
try:
|
||
logger.info(f"Starting incremental sync from {len(self.db_configs)} MySQL database(s) to ChromaDB")
|
||
|
||
# Run all synchronous operations in thread pool
|
||
import asyncio
|
||
loop = asyncio.get_event_loop()
|
||
|
||
def sync_work():
|
||
"""Synchronous work that runs in thread pool"""
|
||
all_chunked_docs = []
|
||
total_docs = 0
|
||
|
||
# Sync from each database
|
||
for db_config in self.db_configs:
|
||
try:
|
||
mysql_sync = self.mysql_syncs[db_config.name]
|
||
last_sync_time = self.last_sync_times.get(db_config.name)
|
||
|
||
# Fetch new documents
|
||
mysql_docs = mysql_sync.fetch_new_documents(last_sync_time)
|
||
|
||
if not mysql_docs:
|
||
logger.debug(f"No new documents in database: {db_config.name}")
|
||
continue
|
||
|
||
# Process and chunk documents
|
||
documents = self.document_processor.process_documents(mysql_docs, db_config)
|
||
chunked_docs = self.document_processor.chunk_documents(documents)
|
||
all_chunked_docs.extend(chunked_docs)
|
||
|
||
# Update synced document IDs for this database
|
||
id_column = db_config.id_column
|
||
new_doc_ids = {str(doc.get(id_column)) for doc in mysql_docs}
|
||
self.synced_doc_ids[db_config.name].update(new_doc_ids)
|
||
self.last_sync_times[db_config.name] = datetime.now()
|
||
total_docs += len(mysql_docs)
|
||
|
||
logger.info(f"Incremental sync: {len(chunked_docs)} chunks from {len(mysql_docs)} documents in database: {db_config.name}")
|
||
except Exception as e:
|
||
logger.error(f"Error during incremental sync for database {db_config.name}: {e}")
|
||
# Continue with other databases
|
||
continue
|
||
|
||
return all_chunked_docs, total_docs
|
||
|
||
# Run synchronous work in thread pool
|
||
all_chunked_docs, total_docs = await loop.run_in_executor(None, sync_work)
|
||
|
||
# Add all new documents to vector store
|
||
if all_chunked_docs:
|
||
# Run in thread pool to avoid blocking event loop
|
||
await loop.run_in_executor(
|
||
None,
|
||
self.vector_store_manager.add_documents,
|
||
all_chunked_docs,
|
||
False # skip_existing，不跳过已存在的文档，默认是更新了内容
|
||
)
|
||
logger.info(f"Incremental sync completed: {len(all_chunked_docs)} chunks from {total_docs} documents across {len(self.db_configs)} database(s)")
|
||
else:
|
||
logger.info("No new documents to sync")
|
||
except Exception as e:
|
||
logger.error(f"Error during incremental sync: {e}")
|
||
raise
|
||
finally:
|
||
self._sync_in_progress = False
|
||
|
||
async def start_auto_sync_with_recovery(self, skip_initial_sync: bool = False):
|
||
"""
|
||
Start automatic periodic sync in background with error recovery.
|
||
If the sync service stops due to an error, it will automatically restart.
|
||
This method runs continuously until stop_auto_sync() is called.
|
||
|
||
Args:
|
||
skip_initial_sync: If True, skip the initial sync_all() call.
|
||
Use this when initial sync is already done elsewhere.
|
||
"""
|
||
if not settings.AUTO_SYNC:
|
||
logger.info("Auto sync is disabled")
|
||
return
|
||
|
||
# Prevent multiple auto sync service instances
|
||
if self._running:
|
||
logger.warning("Auto sync service is already running, skipping duplicate start request")
|
||
return
|
||
|
||
# Check if there's an existing auto sync task still running
|
||
if self._auto_sync_task is not None and not self._auto_sync_task.done():
|
||
logger.warning("Auto sync task is still running, skipping duplicate start request")
|
||
return
|
||
|
||
max_restart_attempts = 10 # Maximum number of restart attempts
|
||
restart_delay = 60 # Wait 60 seconds before restarting after an error
|
||
restart_count = 0
|
||
|
||
self._running = True
|
||
logger.info(f"Auto sync service with error recovery started (interval: {settings.SYNC_INTERVAL}s)")
|
||
|
||
while self._running and restart_count < max_restart_attempts:
|
||
try:
|
||
# Store the task reference to prevent multiple instances
|
||
self._auto_sync_task = asyncio.create_task(
|
||
self._run_auto_sync_loop(skip_initial_sync)
|
||
)
|
||
await self._auto_sync_task
|
||
# If we reach here, the loop exited normally (not due to error)
|
||
logger.info("Auto sync loop exited normally")
|
||
break
|
||
except Exception as e:
|
||
restart_count += 1
|
||
logger.error(
|
||
f"Auto sync service stopped due to error (restart attempt {restart_count}/{max_restart_attempts}): {e}",
|
||
exc_info=True
|
||
)
|
||
|
||
if restart_count >= max_restart_attempts:
|
||
logger.error(f"Auto sync service failed after {max_restart_attempts} restart attempts. Stopping auto sync.")
|
||
self._running = False
|
||
break
|
||
|
||
if not self._running:
|
||
logger.info("Auto sync service stop requested, not restarting")
|
||
break
|
||
|
||
logger.info(f"Waiting {restart_delay}s before restarting auto sync service...")
|
||
await asyncio.sleep(restart_delay)
|
||
logger.info(f"Restarting auto sync service (attempt {restart_count + 1}/{max_restart_attempts})...")
|
||
# Reset skip_initial_sync after first attempt (only skip on first start)
|
||
skip_initial_sync = False
|
||
|
||
logger.info("Auto sync service with recovery stopped")
|
||
self._running = False
|
||
|
||
async def start_auto_sync(self, skip_initial_sync: bool = False):
|
||
"""
|
||
Start automatic periodic sync in background.
|
||
This method runs continuously until stop_auto_sync() is called.
|
||
For production use, prefer start_auto_sync_with_recovery() which includes error recovery.
|
||
|
||
Args:
|
||
skip_initial_sync: If True, skip the initial sync_all() call.
|
||
Use this when initial sync is already done elsewhere.
|
||
"""
|
||
await self._run_auto_sync_loop(skip_initial_sync)
|
||
|
||
async def _run_auto_sync_loop(self, skip_initial_sync: bool = False):
|
||
"""
|
||
Internal method that runs the auto sync loop.
|
||
This is separated so it can be called by both start_auto_sync and start_auto_sync_with_recovery.
|
||
|
||
Args:
|
||
skip_initial_sync: If True, skip the initial sync_all() call.
|
||
"""
|
||
if not settings.AUTO_SYNC:
|
||
logger.info("Auto sync is disabled")
|
||
return
|
||
|
||
logger.info(f"Auto sync loop started (interval: {settings.SYNC_INTERVAL}s)")
|
||
|
||
# Initial sync (only if not skipped)
|
||
if not skip_initial_sync:
|
||
logger.info("Performing initial sync in auto sync service...")
|
||
await self.sync_all()
|
||
else:
|
||
logger.info("Skipping initial sync in auto sync service (already done elsewhere)")
|
||
|
||
# Periodic incremental sync
|
||
sync_count = 0
|
||
last_sync_start_time = None
|
||
while self._running:
|
||
try:
|
||
logger.info(f"Auto sync waiting {settings.SYNC_INTERVAL}s before next sync (count: {sync_count})...")
|
||
await asyncio.sleep(settings.SYNC_INTERVAL)
|
||
if not self._running:
|
||
logger.info("Auto sync stopped, exiting loop")
|
||
break
|
||
|
||
# Check if another sync is in progress (e.g., initial sync or previous incremental sync still running)
|
||
# Wait for it to complete before starting incremental sync (no timeout - wait indefinitely)
|
||
wait_interval = 10 # Check every 10 seconds
|
||
waited_time = 0
|
||
while self._sync_in_progress:
|
||
logger.info(f"Another sync is in progress, waiting... (waited {waited_time}s, will wait until completion)")
|
||
await asyncio.sleep(wait_interval)
|
||
waited_time += wait_interval
|
||
|
||
# Log warning if waiting for a very long time (for monitoring purposes)
|
||
if waited_time % 300 == 0: # Every 5 minutes
|
||
logger.info(f"Still waiting for sync to complete... (waited {waited_time}s / {waited_time // 60} minutes)")
|
||
|
||
# If last sync started more than 2 hours ago and is still running, log warning
|
||
if last_sync_start_time is not None:
|
||
time_since_last_sync = (datetime.now() - last_sync_start_time).total_seconds()
|
||
if time_since_last_sync > 7200: # 2 hours
|
||
logger.warning(f"Last incremental sync has been running for {time_since_last_sync / 3600:.1f} hours, this might indicate a slow sync. Continuing to wait...")
|
||
|
||
if waited_time > 0:
|
||
logger.info(f"Previous sync completed, waited {waited_time}s / {waited_time // 60} minutes")
|
||
|
||
if not self._running:
|
||
logger.info("Auto sync stopped during wait, exiting loop")
|
||
break
|
||
|
||
# Record sync start time for monitoring
|
||
last_sync_start_time = datetime.now()
|
||
sync_count += 1
|
||
logger.info(f"Running periodic incremental sync #{sync_count}...")
|
||
|
||
# Execute incremental sync - it will check _sync_in_progress internally
|
||
try:
|
||
await self.sync_incremental()
|
||
logger.info(f"Incremental sync #{sync_count} completed successfully")
|
||
except Exception as sync_error:
|
||
logger.error(f"Incremental sync #{sync_count} failed: {sync_error}", exc_info=True)
|
||
# Reset sync in progress flag if it was set (in case of unexpected error)
|
||
if self._sync_in_progress:
|
||
logger.warning("Resetting _sync_in_progress flag due to error in incremental sync")
|
||
self._sync_in_progress = False
|
||
# Continue to next cycle even if this sync failed
|
||
logger.info("Continuing to next sync cycle despite error...")
|
||
continue
|
||
|
||
except asyncio.CancelledError:
|
||
logger.info("Auto sync task was cancelled")
|
||
break
|
||
except Exception as e:
|
||
logger.error(f"Unexpected error in auto sync loop: {e}", exc_info=True)
|
||
# Reset sync in progress flag if it was set (in case of unexpected error)
|
||
if self._sync_in_progress:
|
||
logger.warning("Resetting _sync_in_progress flag due to unexpected error in auto sync loop")
|
||
self._sync_in_progress = False
|
||
# Continue running even if there's an error
|
||
logger.info("Continuing auto sync loop despite error...")
|
||
continue
|
||
|
||
logger.info(f"Auto sync loop stopped (total syncs: {sync_count})")
|
||
|
||
def stop_auto_sync(self):
|
||
"""Stop automatic sync"""
|
||
if not self._running:
|
||
logger.info("Auto sync service is not running")
|
||
return
|
||
|
||
self._running = False
|
||
logger.info("Auto sync service stop requested")
|
||
|
||
# Wait for current sync to complete (if any) before stopping
|
||
# This prevents stopping in the middle of a sync operation
|
||
if self._sync_in_progress:
|
||
logger.info("Waiting for current sync to complete before stopping auto sync service...")
|
||
import time
|
||
max_wait = 300 # Wait up to 5 minutes
|
||
waited = 0
|
||
while self._sync_in_progress and waited < max_wait:
|
||
time.sleep(2)
|
||
waited += 2
|
||
if self._sync_in_progress:
|
||
logger.warning(f"Sync still in progress after waiting {waited}s, forcing stop")
|
||
else:
|
||
logger.info("Current sync completed, auto sync service stopped")
|
||
else:
|
||
logger.info("Auto sync service stopped")
|
||
|
||
def close(self):
|
||
"""Close MySQL connections"""
|
||
self.stop_auto_sync()
|
||
for mysql_sync in self.mysql_syncs.values():
|
||
if hasattr(mysql_sync, 'close'):
|
||
mysql_sync.close()
|
||
|