339 lines
11 KiB
Python
339 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script for testing document sync with various database configurations
|
|
"""
|
|
import os
|
|
import sys
|
|
import json
|
|
import asyncio
|
|
from datetime import datetime
|
|
import argparse
|
|
|
|
# Add current directory to Python path
|
|
sys.path.append('.')
|
|
|
|
from config import Settings, DatabaseDataSourceConfig
|
|
from sync_service import SyncService
|
|
|
|
|
|
class ConfigurationLoader:
|
|
"""Load configurations from database_config.json"""
|
|
|
|
def __init__(self, config_file_path: str = "./database_config.json"):
|
|
"""
|
|
Initialize configuration loader
|
|
|
|
Args:
|
|
config_file_path: Path to the database configuration file
|
|
"""
|
|
self.config_file_path = config_file_path
|
|
self.configurations = []
|
|
|
|
# Load configurations from file
|
|
self._load_configurations()
|
|
|
|
def _load_configurations(self):
|
|
"""Load configurations from the database_config.json file"""
|
|
try:
|
|
with open(self.config_file_path, "r", encoding="utf-8") as f:
|
|
self.configurations = json.load(f)
|
|
|
|
print(f"Loaded {len(self.configurations)} configurations from {self.config_file_path}")
|
|
|
|
# Print summary of available configurations
|
|
print("\nAvailable configurations:")
|
|
for i, config in enumerate(self.configurations, 1):
|
|
print(f"{i}. {config['name']} (database: {config['database']}, table: {config['table_name']})")
|
|
print(f" Content columns: {config['content_column']}")
|
|
if config.get('file_source_type'):
|
|
print(f" File source: {config['file_source_type']}")
|
|
print()
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to load configurations from {self.config_file_path}: {e}")
|
|
|
|
def get_configuration_by_name(self, config_name: str):
|
|
"""
|
|
Get configuration by name
|
|
|
|
Args:
|
|
config_name: Name of the configuration to retrieve
|
|
|
|
Returns:
|
|
Configuration dictionary
|
|
"""
|
|
for config in self.configurations:
|
|
if config['name'] == config_name:
|
|
return config
|
|
raise ValueError(f"Configuration with name '{config_name}' not found")
|
|
|
|
def get_configurations_by_type(self, config_type: str = "database"):
|
|
"""
|
|
Get configurations by type
|
|
|
|
Args:
|
|
config_type: Type of configurations to retrieve
|
|
|
|
Returns:
|
|
List of configuration dictionaries
|
|
"""
|
|
return [config for config in self.configurations if config['type'] == config_type]
|
|
|
|
def get_all_configurations(self):
|
|
"""
|
|
Get all configurations
|
|
|
|
Returns:
|
|
List of all configuration dictionaries
|
|
"""
|
|
return self.configurations
|
|
|
|
def create_test_config_file(self, configurations):
|
|
"""
|
|
Create a temporary test configuration file
|
|
|
|
Args:
|
|
configurations: List of configurations to include in the test file
|
|
|
|
Returns:
|
|
Path to the created test configuration file
|
|
"""
|
|
config_path = f"./test_config_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
|
|
|
|
with open(config_path, "w", encoding="utf-8") as f:
|
|
json.dump(configurations, f, ensure_ascii=False, indent=2)
|
|
|
|
print(f"Created test configuration file: {config_path}")
|
|
return config_path
|
|
|
|
|
|
async def test_sync(configurations, sync_type: str = "all", force: bool = False):
|
|
"""
|
|
Test document sync with the specified configurations
|
|
|
|
Args:
|
|
configurations: List of configurations to test
|
|
sync_type: Type of sync to perform ("all" for full sync, "incremental" for incremental sync)
|
|
force: Whether to force full sync (ignored for incremental)
|
|
"""
|
|
if not configurations:
|
|
print("No configurations to test")
|
|
return False
|
|
|
|
config_names = [config["name"] for config in configurations]
|
|
print(f"\n=== Testing {sync_type.upper()} Sync with Configuration(s): {', '.join(config_names)} ===")
|
|
|
|
loader = ConfigurationLoader()
|
|
|
|
# Define database path outside try block for cleanup
|
|
import sqlite3
|
|
from pathlib import Path
|
|
import json
|
|
|
|
DATA_DIR = Path(__file__).parent / "data"
|
|
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
|
DB_PATH = DATA_DIR / "sessions.db"
|
|
|
|
try:
|
|
# Write test configurations to SQLite database
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Create table if not exists
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS data_sources (
|
|
name TEXT PRIMARY KEY,
|
|
config TEXT,
|
|
update_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# Insert/update test configurations
|
|
for config in configurations:
|
|
config_json = json.dumps(config, ensure_ascii=False, indent=2)
|
|
cursor.execute('''
|
|
INSERT OR REPLACE INTO data_sources (name, config, update_at)
|
|
VALUES (?, ?, CURRENT_TIMESTAMP)
|
|
''', (config['name'], config_json))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
print(f"Inserted {len(configurations)} test configurations into SQLite database")
|
|
|
|
# Initialize settings normally (will read from SQLite)
|
|
from config import Settings
|
|
settings = Settings()
|
|
|
|
# Initialize sync service
|
|
print("\nInitializing SyncService...")
|
|
sync_service = SyncService()
|
|
|
|
print(f"\nConfiguration details:")
|
|
for db_config in sync_service.db_configs:
|
|
print(f"- Name: {db_config.name}")
|
|
print(f"- Database: {db_config.database}")
|
|
print(f"- Table: {db_config.table_name}")
|
|
print(f"- Content columns: {db_config.content_columns}")
|
|
if db_config.file_source_type:
|
|
print(f"- File source: {db_config.file_source_type}")
|
|
if db_config.file_system_base_path:
|
|
print(f"- File base path: {db_config.file_system_base_path}")
|
|
if db_config.updated_at_column:
|
|
print(f"- Updated at column: {db_config.updated_at_column}")
|
|
print()
|
|
|
|
# Perform sync
|
|
start_time = datetime.now()
|
|
|
|
if sync_type == "all":
|
|
print(f"Starting full sync (force={force})...")
|
|
await sync_service.sync_all(force=force)
|
|
elif sync_type == "incremental":
|
|
print("Starting incremental sync...")
|
|
await sync_service.sync_incremental()
|
|
else:
|
|
raise ValueError(f"Unknown sync type: {sync_type}")
|
|
|
|
duration = (datetime.now() - start_time).total_seconds()
|
|
print(f"\n✓ Sync completed successfully in {duration:.2f} seconds")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"\n✗ Sync failed: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
finally:
|
|
# Clean up: remove test configurations from SQLite database
|
|
try:
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Delete test configurations
|
|
for config in configurations:
|
|
cursor.execute('DELETE FROM data_sources WHERE name = ?', (config['name'],))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
print(f"Removed {len(configurations)} test configurations from SQLite database")
|
|
except Exception as e:
|
|
print(f"Error during cleanup: {e}")
|
|
|
|
|
|
async def main():
|
|
"""
|
|
Main function to run tests
|
|
"""
|
|
parser = argparse.ArgumentParser(description="Test document sync with configurations from database_config.json")
|
|
|
|
# Add mutually exclusive group for configuration selection
|
|
group = parser.add_mutually_exclusive_group()
|
|
|
|
group.add_argument(
|
|
"--all",
|
|
action="store_true",
|
|
help="Test all configurations"
|
|
)
|
|
|
|
group.add_argument(
|
|
"--name",
|
|
type=str,
|
|
help="Test specific configuration by name"
|
|
)
|
|
|
|
group.add_argument(
|
|
"--list",
|
|
action="store_true",
|
|
help="List all available configurations and exit"
|
|
)
|
|
|
|
# Sync options
|
|
parser.add_argument(
|
|
"--sync-type",
|
|
type=str,
|
|
choices=["all", "incremental"],
|
|
default="all",
|
|
help="Type of sync to perform (default: all)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--force",
|
|
action="store_true",
|
|
help="Force full sync (ignored for incremental sync)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
print("=" * 70)
|
|
print("RAG Sync Configuration Tester")
|
|
print("=" * 70)
|
|
|
|
# Initialize configuration loader
|
|
loader = ConfigurationLoader()
|
|
|
|
# List configurations and exit if requested
|
|
if args.list:
|
|
print("\nConfiguration listing complete.")
|
|
return 0
|
|
|
|
# Determine which configurations to test
|
|
configurations_to_test = []
|
|
|
|
if args.all:
|
|
# Test all configurations
|
|
configurations_to_test = loader.get_all_configurations()
|
|
print(f"\nSelected all {len(configurations_to_test)} configurations for testing")
|
|
|
|
elif args.name:
|
|
# Test specific configuration by name
|
|
try:
|
|
config = loader.get_configuration_by_name(args.name)
|
|
configurations_to_test = [config]
|
|
print(f"\nSelected configuration: {args.name}")
|
|
except ValueError as e:
|
|
print(f"\nError: {e}")
|
|
return 1
|
|
|
|
else:
|
|
# Default: ask user to select configuration
|
|
print("\nPlease select configuration(s) to test (comma-separated numbers or 'all'):")
|
|
print("Example: 1,3 or all")
|
|
|
|
user_input = input("Selection: ").strip()
|
|
|
|
if user_input.lower() == "all":
|
|
configurations_to_test = loader.get_all_configurations()
|
|
else:
|
|
try:
|
|
indices = [int(idx.strip()) - 1 for idx in user_input.split(",")]
|
|
configurations_to_test = [loader.configurations[i] for i in indices]
|
|
except (ValueError, IndexError) as e:
|
|
print(f"\nInvalid selection: {e}")
|
|
return 1
|
|
|
|
# Test the selected configurations
|
|
if configurations_to_test:
|
|
print(f"\nTesting {len(configurations_to_test)} configuration(s)")
|
|
|
|
# Test configurations
|
|
success = await test_sync(configurations_to_test, args.sync_type, args.force)
|
|
|
|
# Print results
|
|
print("\n" + "=" * 70)
|
|
if success:
|
|
print("🎉 All selected configurations passed the sync test!")
|
|
return 0
|
|
else:
|
|
print("❌ Some configurations failed the sync test.")
|
|
return 1
|
|
else:
|
|
print("\nNo configurations selected for testing.")
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(asyncio.run(main()))
|