287 lines
10 KiB
Python
287 lines
10 KiB
Python
"""Remote folder synchronization implementation"""
|
||
import os
|
||
import re
|
||
from typing import List, Dict, Any, Set
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from loguru import logger
|
||
from config import BaseDataSourceConfig
|
||
from sync.base_sync import BaseSync
|
||
from rag.file_parser import FileParser
|
||
|
||
|
||
class RemoteFolderSync(BaseSync):
|
||
"""Handle synchronization between remote folder (via SCP) and ChromaDB"""
|
||
|
||
def __init__(self, config: BaseDataSourceConfig):
|
||
"""
|
||
Initialize remote folder sync with configuration
|
||
|
||
Args:
|
||
config: Remote folder configuration
|
||
"""
|
||
super().__init__(config)
|
||
self.file_parser = FileParser()
|
||
self._ssh_client = None
|
||
self._sftp_client = None
|
||
|
||
def fetch_all_documents(self) -> List[Dict[str, Any]]:
|
||
"""
|
||
Fetch all documents from the remote folder
|
||
|
||
Returns:
|
||
List of documents
|
||
"""
|
||
return self._fetch_documents()
|
||
|
||
def fetch_new_documents(self, last_sync_time=None, synced_doc_ids=None) -> List[Dict[str, Any]]:
|
||
"""
|
||
Fetch new/updated documents from the remote folder since last sync time
|
||
|
||
Args:
|
||
last_sync_time: Last synchronization time
|
||
synced_doc_ids: Set of document IDs that have already been synced
|
||
|
||
Returns:
|
||
List of new/updated documents
|
||
"""
|
||
return self._fetch_documents(last_sync_time, synced_doc_ids)
|
||
|
||
def get_synced_document_ids(self) -> Set[str]:
|
||
"""
|
||
Get IDs of all files in the remote folder
|
||
|
||
Returns:
|
||
Set of file paths (as document IDs)
|
||
"""
|
||
files = self._get_all_files()
|
||
return {str(file) for file in files}
|
||
|
||
def _fetch_documents(self, last_sync_time=None, synced_doc_ids=None) -> List[Dict[str, Any]]:
|
||
"""
|
||
Internal method to fetch documents from remote folder via SFTP
|
||
|
||
Args:
|
||
last_sync_time: Last synchronization time (for incremental sync)
|
||
synced_doc_ids: Set of document IDs that have already been synced
|
||
|
||
Returns:
|
||
List of documents
|
||
"""
|
||
documents = []
|
||
self._connect()
|
||
try:
|
||
files = self._get_all_files()
|
||
|
||
for file_path in files:
|
||
# Check if document has already been synced
|
||
if synced_doc_ids and file_path in synced_doc_ids:
|
||
# If file is already synced, check if it's been modified since last sync
|
||
if last_sync_time:
|
||
file_stat = self._sftp_client.stat(file_path)
|
||
file_mtime = datetime.fromtimestamp(file_stat.st_mtime)
|
||
# Skip if not modified since last sync
|
||
if file_mtime <= last_sync_time:
|
||
continue
|
||
# If file is not synced yet, always include it regardless of modification time
|
||
# This handles the case where files were added to the folder after last_sync_time but have older mtimes
|
||
|
||
# Parse file content
|
||
try:
|
||
# 使用 BaseSync 中的通用方法获取文件字节
|
||
# 对于远程文件，我们需要先通过 SFTP 获取文件内容
|
||
with self._sftp_client.open(file_path, 'rb') as f:
|
||
file_bytes = f.read()
|
||
|
||
# 检查文件扩展名是否在支持的列表中
|
||
if Path(file_path).suffix.lower() not in FileParser.SUPPORTED_EXTENSIONS:
|
||
logger.debug(f"Skipping unsupported file: {file_path}")
|
||
continue
|
||
|
||
# 使用 BaseSync 中的通用方法解析文件内容，传递远程主机地址
|
||
content = self._parse_file_content(file_bytes, file_path, host=self.config.host) if file_bytes else f"[无法读取文件：{Path(file_path).name}]"
|
||
|
||
document = {
|
||
'id': str(file_path),
|
||
'title': Path(file_path).name,
|
||
'content': content,
|
||
'file_path': str(file_path),
|
||
'update_time': datetime.fromtimestamp(self._sftp_client.stat(file_path).st_mtime)
|
||
}
|
||
documents.append(document)
|
||
except Exception as e:
|
||
logger.error(f"Error processing file {file_path}: {e}")
|
||
finally:
|
||
self._disconnect()
|
||
|
||
return documents
|
||
|
||
def _get_all_files(self) -> List[str]:
|
||
"""
|
||
Get all files in the remote folder
|
||
|
||
Returns:
|
||
List of file paths
|
||
"""
|
||
files = []
|
||
self._get_files_recursive(self.config.folder_path, files)
|
||
return files
|
||
|
||
def _get_files_recursive(self, folder_path: str, files: List[str]):
|
||
"""
|
||
Recursively get all files in the remote folder
|
||
|
||
Args:
|
||
folder_path: Current folder path
|
||
files: List to store found files
|
||
"""
|
||
try:
|
||
items = self._sftp_client.listdir_attr(folder_path)
|
||
|
||
for item in items:
|
||
item_path = os.path.join(folder_path, item.filename)
|
||
|
||
if item.filename not in ('.', '..'):
|
||
if item.st_mode & 0o040000: # Check if it's a directory
|
||
if self.config.recursive:
|
||
self._get_files_recursive(item_path, files)
|
||
else:
|
||
# Check if file should be ignored
|
||
if not self._should_ignore_file(item_path):
|
||
files.append(item_path)
|
||
except Exception as e:
|
||
logger.error(f"Error listing remote folder {folder_path}: {e}")
|
||
|
||
def _should_ignore_file(self, file_path: str) -> bool:
|
||
"""
|
||
Check if file should be ignored based on ignore patterns
|
||
|
||
Args:
|
||
file_path: File path to check
|
||
|
||
Returns:
|
||
True if file should be ignored, False otherwise
|
||
"""
|
||
if not hasattr(self.config, 'ignore_patterns') or not self.config.ignore_patterns:
|
||
return False
|
||
|
||
# Get relative path from folder root
|
||
relative_path = os.path.relpath(file_path, self.config.folder_path)
|
||
|
||
for pattern in self.config.ignore_patterns:
|
||
if self._match_pattern(relative_path, pattern):
|
||
return True
|
||
|
||
return False
|
||
|
||
def _match_pattern(self, path: str, pattern: str) -> bool:
|
||
"""
|
||
Match a path against a pattern (similar to .gitignore)
|
||
|
||
Args:
|
||
path: Path to match
|
||
pattern: Pattern to match against
|
||
|
||
Returns:
|
||
True if path matches pattern, False otherwise
|
||
"""
|
||
# Convert glob pattern to regex
|
||
regex_pattern = pattern
|
||
regex_pattern = regex_pattern.replace('.', r'\.')
|
||
regex_pattern = regex_pattern.replace('*', r'.*')
|
||
regex_pattern = regex_pattern.replace('?', r'.')
|
||
|
||
# Handle directory patterns
|
||
if pattern.endswith('/'):
|
||
regex_pattern = f'^{regex_pattern}.*$'
|
||
else:
|
||
regex_pattern = f'^{regex_pattern}$'
|
||
|
||
return bool(re.match(regex_pattern, path))
|
||
|
||
def _connect(self):
|
||
"""
|
||
Connect to the remote server via SSH/SFTP
|
||
"""
|
||
import paramiko
|
||
|
||
self._ssh_client = paramiko.SSHClient()
|
||
self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
|
||
# Connect to SSH server
|
||
# 确保username不为None，否则paramiko会报错
|
||
username = self.config.username or ''
|
||
|
||
# 使用密码认证或SSH代理认证
|
||
self._ssh_client.connect(
|
||
hostname=self.config.host,
|
||
port=self.config.port or 22,
|
||
username=username,
|
||
password=self.config.password,
|
||
timeout=10,
|
||
allow_agent=True, # 允许使用SSH代理
|
||
look_for_keys=False # 禁用查找本地密钥文件
|
||
)
|
||
|
||
# Create SFTP client
|
||
self._sftp_client = self._ssh_client.open_sftp()
|
||
|
||
def _disconnect(self):
|
||
"""
|
||
Disconnect from the remote server
|
||
"""
|
||
if self._sftp_client:
|
||
self._sftp_client.close()
|
||
self._sftp_client = None
|
||
|
||
if self._ssh_client:
|
||
self._ssh_client.close()
|
||
self._ssh_client = None
|
||
|
||
@staticmethod
|
||
def check_data_source_exists(config: BaseDataSourceConfig) -> bool:
|
||
"""
|
||
Check if the remote folder exists and is accessible
|
||
|
||
Args:
|
||
config: Remote folder configuration
|
||
|
||
Returns:
|
||
True if remote folder exists and is accessible, False otherwise
|
||
"""
|
||
import paramiko
|
||
|
||
ssh_client = None
|
||
sftp_client = None
|
||
|
||
try:
|
||
ssh_client = paramiko.SSHClient()
|
||
ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
|
||
# Connect to SSH server - use SSH agent if available, otherwise password
|
||
username = config.username or ''
|
||
ssh_client.connect(
|
||
hostname=config.host,
|
||
port=config.port or 22,
|
||
username=username,
|
||
password=config.password,
|
||
timeout=10,
|
||
allow_agent=True, # 允许使用SSH代理
|
||
look_for_keys=False # 禁用查找本地密钥文件
|
||
)
|
||
|
||
# Create SFTP client and check folder exists
|
||
sftp_client = ssh_client.open_sftp()
|
||
sftp_client.stat(config.folder_path)
|
||
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"Error checking remote folder: {e}")
|
||
return False
|
||
finally:
|
||
if sftp_client:
|
||
sftp_client.close()
|
||
if ssh_client:
|
||
ssh_client.close()
|
||
|