290 lines
10 KiB
Python
290 lines
10 KiB
Python
"""Dameng database synchronization implementation"""
|
|
import urllib
|
|
import dmPython # 达梦数据库驱动
|
|
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 DaMengSync(BaseSync):
|
|
"""Handle synchronization between Dameng database and ChromaDB"""
|
|
|
|
def __init__(self, db_config: DatabaseConfig, vector_store_manager=None):
|
|
"""
|
|
Initialize DaMeng 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 Dameng database 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
|
|
# database = self.db_config.database or settings.DATABASE
|
|
|
|
# Log connection info (without password)
|
|
logger.info(f"Connecting to DaMeng Schema: {user}")
|
|
logger.debug(f"DaMeng connection details: host={host}, port={port}, user={user}")
|
|
|
|
# Create connection
|
|
self.connection = dmPython.connect(
|
|
user = user,
|
|
password = password,
|
|
server = host,
|
|
port = port
|
|
)
|
|
|
|
def fetch_all_documents(self, last_sync_time=None) -> List[Dict[str, Any]]:
|
|
"""
|
|
Fetch documents from the DaMeng database
|
|
|
|
Args:
|
|
last_sync_time: Last synchronization time (for incremental sync)
|
|
|
|
Returns:
|
|
List of documents
|
|
"""
|
|
self._connect()
|
|
cursor = self.connection.cursor()
|
|
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} > ?"
|
|
params.append(last_sync_time)
|
|
|
|
logger.debug(f"DaMeng query: {query}, params: {params}")
|
|
cursor.execute(query, params)
|
|
|
|
# Get column names from cursor description
|
|
column_names = [desc[0] for desc in cursor.description]
|
|
|
|
# Parse results
|
|
documents = []
|
|
for row in cursor.fetchall():
|
|
# Convert row tuple to dict
|
|
row_dict = dict(zip(column_names, row))
|
|
|
|
# Handle file content if file_column is specified
|
|
if self.db_config.file_column and row_dict.get(self.db_config.file_column):
|
|
# Extract file path from the file column
|
|
file_path = row_dict[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_dict[self.db_config.content_column] = file_content
|
|
|
|
# Generate unique document ID
|
|
record_id = str(row_dict[self.db_config.id_column])
|
|
doc_id = self.generate_doc_id(record_id)
|
|
row_dict['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_dict.get(
|
|
self.db_config.updated_at_column):
|
|
# Skip if not modified since last sync
|
|
if row_dict[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_dict)
|
|
|
|
return documents
|
|
finally:
|
|
cursor.close()
|
|
|
|
def fetch_new_documents(self, last_sync_time=None) -> List[Dict[str, Any]]:
|
|
"""
|
|
Fetch new/updated documents from the Dameng 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 Dameng database
|
|
|
|
Returns:
|
|
Set of document IDs
|
|
"""
|
|
cursor = self.connection.cursor()
|
|
try:
|
|
query = f"SELECT {self.db_config.id_column} FROM {self.db_config.table_name}"
|
|
cursor.execute(query)
|
|
column_names = [desc[0] for desc in cursor.description]
|
|
id_index = column_names.index(self.db_config.id_column)
|
|
return {str(row[id_index]) for row in cursor.fetchall()}
|
|
finally:
|
|
cursor.close()
|
|
|
|
def generate_doc_id(self, record_id: str) -> str:
|
|
"""
|
|
Generate a unique document ID for Dameng records
|
|
|
|
Args:
|
|
record_id: ID of the record in the database
|
|
|
|
Returns:
|
|
Unique document ID based on user, table, and record ID
|
|
"""
|
|
# 为达梦数据库记录生成唯一的文档 ID
|
|
return f"{self.db_config.user}_dameng_{self.db_config.table_name}_{record_id}"
|
|
|
|
def doc_to_llamaindex_doc(self, doc: Dict) -> 'Document':
|
|
"""
|
|
Convert Dameng document to LlamaIndex Document
|
|
|
|
Args:
|
|
doc: Dameng 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": "dameng",
|
|
# Dameng中没有database的概念
|
|
"schema": self.db_config.user,
|
|
"table": self.db_config.table_name,
|
|
"db_source": self.db_config.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 DaMeng table exists and is accessible
|
|
|
|
Args:
|
|
config: Database configuration
|
|
|
|
Returns:
|
|
True if table exists and is accessible, False otherwise
|
|
"""
|
|
temp_connection = None
|
|
try:
|
|
# Get connection parameters
|
|
host = config.host or settings.HOST
|
|
port = config.port or settings.PORT
|
|
user = config.user or settings.USER
|
|
password = config.password or settings.PASSWORD
|
|
|
|
# Create connection
|
|
temp_connection = dmPython.connect(
|
|
user=user,
|
|
password=password,
|
|
server=host,
|
|
port=port
|
|
)
|
|
|
|
# Check if table exists
|
|
cursor = temp_connection.cursor()
|
|
cursor.execute(
|
|
# "SELECT TABLE_NAME FROM USER_TABLES WHERE TABLE_NAME = ?",
|
|
"SELECT COUNT(*) FROM ALL_TABLES WHERE OWNER = ? AND TABLE_NAME = ?",
|
|
[config.user.upper(), config.table_name.upper()] # DM 模式名、表名默认大写
|
|
)
|
|
result = cursor.fetchone()
|
|
cursor.close()
|
|
|
|
# 判断查询结果
|
|
if result and result[0] > 0:
|
|
return True # 表存在
|
|
else:
|
|
logger.error(f"Table {config.table_name} does not exist in database {config.database}")
|
|
return False # 表不存在
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error checking Dameng data source: {e}")
|
|
return False
|
|
finally:
|
|
if temp_connection:
|
|
temp_connection.close()
|