mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-27 02:42:00 +08:00
360 lines
13 KiB
Python
360 lines
13 KiB
Python
"""
|
|
RAGFlow-specific migration orchestrator from Elasticsearch to OceanBase.
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
from typing import Any, Callable
|
|
|
|
from rich.console import Console
|
|
from rich.progress import (
|
|
Progress,
|
|
SpinnerColumn,
|
|
TextColumn,
|
|
BarColumn,
|
|
TaskProgressColumn,
|
|
TimeRemainingColumn,
|
|
)
|
|
|
|
from .es_client import ESClient
|
|
from .ob_client import OBClient
|
|
from .schema import RAGFlowSchemaConverter, RAGFlowDataConverter
|
|
from .progress import ProgressManager, MigrationProgress
|
|
from .verify import MigrationVerifier
|
|
|
|
logger = logging.getLogger(__name__)
|
|
console = Console()
|
|
|
|
|
|
class ESToOceanBaseMigrator:
|
|
"""
|
|
RAGFlow-specific migration orchestrator.
|
|
|
|
This migrator is designed specifically for RAGFlow's data structure,
|
|
handling the fixed schema and vector embeddings correctly.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
es_client: ESClient,
|
|
ob_client: OBClient,
|
|
progress_dir: str = ".migration_progress",
|
|
):
|
|
"""
|
|
Initialize migrator.
|
|
|
|
Args:
|
|
es_client: Elasticsearch client
|
|
ob_client: OceanBase client
|
|
progress_dir: Directory for progress files
|
|
"""
|
|
self.es_client = es_client
|
|
self.ob_client = ob_client
|
|
self.progress_manager = ProgressManager(progress_dir)
|
|
self.schema_converter = RAGFlowSchemaConverter()
|
|
|
|
def migrate(
|
|
self,
|
|
es_index: str,
|
|
ob_table: str,
|
|
batch_size: int = 1000,
|
|
resume: bool = False,
|
|
verify_after: bool = True,
|
|
on_progress: Callable[[int, int], None] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Execute full migration from ES to OceanBase for RAGFlow data.
|
|
|
|
Args:
|
|
es_index: Source Elasticsearch index
|
|
ob_table: Target OceanBase table
|
|
batch_size: Documents per batch
|
|
resume: Resume from previous progress
|
|
verify_after: Run verification after migration
|
|
on_progress: Progress callback (migrated, total)
|
|
|
|
Returns:
|
|
Migration result dictionary
|
|
"""
|
|
start_time = time.time()
|
|
result = {
|
|
"success": False,
|
|
"es_index": es_index,
|
|
"ob_table": ob_table,
|
|
"total_documents": 0,
|
|
"migrated_documents": 0,
|
|
"failed_documents": 0,
|
|
"duration_seconds": 0,
|
|
"verification": None,
|
|
"error": None,
|
|
}
|
|
|
|
progress: MigrationProgress | None = None
|
|
|
|
try:
|
|
# Step 1: Check connections
|
|
console.print("[bold blue]Step 1: Checking connections...[/]")
|
|
self._check_connections()
|
|
|
|
# Step 2: Analyze ES index
|
|
console.print("\n[bold blue]Step 2: Analyzing ES index...[/]")
|
|
analysis = self._analyze_es_index(es_index)
|
|
|
|
# Auto-detect vector size from ES mapping
|
|
vector_size = 768 # Default fallback
|
|
if analysis["vector_fields"]:
|
|
vector_size = analysis["vector_fields"][0]["dimension"]
|
|
console.print(f" [green]Auto-detected vector dimension: {vector_size}[/]")
|
|
else:
|
|
console.print(f" [yellow]No vector fields found, using default: {vector_size}[/]")
|
|
console.print(f" Known RAGFlow fields: {len(analysis['known_fields'])}")
|
|
if analysis["unknown_fields"]:
|
|
console.print(f" [yellow]Unknown fields (will be stored in 'extra'): {analysis['unknown_fields']}[/]")
|
|
|
|
# Step 3: Get total document count
|
|
total_docs = self.es_client.count_documents(es_index)
|
|
console.print(f" Total documents: {total_docs:,}")
|
|
|
|
result["total_documents"] = total_docs
|
|
|
|
if total_docs == 0:
|
|
console.print("[yellow]No documents to migrate[/]")
|
|
result["success"] = True
|
|
return result
|
|
|
|
# Step 4: Handle resume or fresh start
|
|
if resume and self.progress_manager.can_resume(es_index, ob_table):
|
|
console.print("\n[bold yellow]Resuming from previous progress...[/]")
|
|
progress = self.progress_manager.load_progress(es_index, ob_table)
|
|
console.print(f" Previously migrated: {progress.migrated_documents:,} documents")
|
|
else:
|
|
# Fresh start - check if table already exists
|
|
if self.ob_client.table_exists(ob_table):
|
|
raise RuntimeError(f"Table '{ob_table}' already exists in OceanBase. Migration aborted to prevent data conflicts. Please drop the table manually or use a different table name.")
|
|
|
|
progress = self.progress_manager.create_progress(es_index, ob_table, total_docs)
|
|
|
|
# Step 5: Create table if needed
|
|
if not progress.table_created:
|
|
console.print("\n[bold blue]Step 3: Creating OceanBase table...[/]")
|
|
if not self.ob_client.table_exists(ob_table):
|
|
self.ob_client.create_ragflow_table(
|
|
table_name=ob_table,
|
|
vector_size=vector_size,
|
|
create_indexes=True,
|
|
create_fts_indexes=True,
|
|
)
|
|
console.print(f" Created table '{ob_table}' with RAGFlow schema")
|
|
else:
|
|
console.print(f" Table '{ob_table}' already exists")
|
|
# Check and add vector column if needed
|
|
self.ob_client.add_vector_column(ob_table, vector_size)
|
|
|
|
progress.table_created = True
|
|
progress.indexes_created = True
|
|
progress.schema_converted = True
|
|
self.progress_manager.save_progress(progress)
|
|
|
|
# Step 6: Migrate data
|
|
console.print("\n[bold blue]Step 4: Migrating data...[/]")
|
|
data_converter = RAGFlowDataConverter()
|
|
|
|
migrated = self._migrate_data(
|
|
es_index=es_index,
|
|
ob_table=ob_table,
|
|
data_converter=data_converter,
|
|
progress=progress,
|
|
batch_size=batch_size,
|
|
on_progress=on_progress,
|
|
)
|
|
|
|
result["migrated_documents"] = migrated
|
|
result["failed_documents"] = progress.failed_documents
|
|
|
|
# Step 7: Mark completed
|
|
self.progress_manager.mark_completed(progress)
|
|
|
|
# Step 8: Verify (optional)
|
|
if verify_after:
|
|
console.print("\n[bold blue]Step 5: Verifying migration...[/]")
|
|
verifier = MigrationVerifier(self.es_client, self.ob_client)
|
|
verification = verifier.verify(es_index, ob_table, primary_key="id")
|
|
result["verification"] = {
|
|
"passed": verification.passed,
|
|
"message": verification.message,
|
|
"es_count": verification.es_count,
|
|
"ob_count": verification.ob_count,
|
|
"sample_match_rate": verification.sample_match_rate,
|
|
}
|
|
console.print(verifier.generate_report(verification))
|
|
|
|
result["success"] = True
|
|
result["duration_seconds"] = time.time() - start_time
|
|
|
|
console.print(
|
|
f"\n[bold green]Migration completed successfully![/]"
|
|
f"\n Total: {result['total_documents']:,} documents"
|
|
f"\n Migrated: {result['migrated_documents']:,} documents"
|
|
f"\n Failed: {result['failed_documents']:,} documents"
|
|
f"\n Duration: {result['duration_seconds']:.1f} seconds"
|
|
)
|
|
|
|
except KeyboardInterrupt:
|
|
console.print("\n[bold yellow]Migration interrupted by user[/]")
|
|
if progress:
|
|
self.progress_manager.mark_paused(progress)
|
|
result["error"] = "Interrupted by user"
|
|
|
|
except Exception as e:
|
|
logger.exception("Migration failed")
|
|
if progress:
|
|
self.progress_manager.mark_failed(progress, str(e))
|
|
result["error"] = str(e)
|
|
console.print(f"\n[bold red]Migration failed: {e}[/]")
|
|
|
|
return result
|
|
|
|
def _check_connections(self):
|
|
"""Verify connections to both databases."""
|
|
# Check ES
|
|
es_health = self.es_client.health_check()
|
|
if es_health.get("status") not in ("green", "yellow"):
|
|
raise RuntimeError(f"ES cluster unhealthy: {es_health}")
|
|
console.print(f" ES cluster status: {es_health.get('status')}")
|
|
|
|
# Check OceanBase
|
|
if not self.ob_client.health_check():
|
|
raise RuntimeError("OceanBase connection failed")
|
|
|
|
ob_version = self.ob_client.get_version()
|
|
console.print(f" OceanBase connection: OK (version: {ob_version})")
|
|
|
|
def _analyze_es_index(self, es_index: str) -> dict[str, Any]:
|
|
"""Analyze ES index structure for RAGFlow compatibility."""
|
|
es_mapping = self.es_client.get_index_mapping(es_index)
|
|
return self.schema_converter.analyze_es_mapping(es_mapping)
|
|
|
|
def _migrate_data(
|
|
self,
|
|
es_index: str,
|
|
ob_table: str,
|
|
data_converter: RAGFlowDataConverter,
|
|
progress: MigrationProgress,
|
|
batch_size: int,
|
|
on_progress: Callable[[int, int], None] | None,
|
|
) -> int:
|
|
"""Migrate data in batches."""
|
|
total = progress.total_documents
|
|
migrated = progress.migrated_documents
|
|
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TaskProgressColumn(),
|
|
TimeRemainingColumn(),
|
|
console=console,
|
|
) as pbar:
|
|
task = pbar.add_task(
|
|
"Migrating...",
|
|
total=total,
|
|
completed=migrated,
|
|
)
|
|
|
|
batch_count = 0
|
|
for batch in self.es_client.scroll_documents(es_index, batch_size):
|
|
batch_count += 1
|
|
|
|
# Convert batch to OceanBase format
|
|
ob_rows = data_converter.convert_batch(batch)
|
|
|
|
# Insert batch
|
|
try:
|
|
inserted = self.ob_client.insert_batch(ob_table, ob_rows)
|
|
migrated += inserted
|
|
|
|
# Update progress
|
|
last_ids = [doc.get("_id", doc.get("id", "")) for doc in batch]
|
|
self.progress_manager.update_progress(
|
|
progress,
|
|
migrated_count=inserted,
|
|
last_batch_ids=last_ids,
|
|
)
|
|
|
|
# Update progress bar
|
|
pbar.update(task, completed=migrated)
|
|
|
|
# Callback
|
|
if on_progress:
|
|
on_progress(migrated, total)
|
|
|
|
# Log periodically
|
|
if batch_count % 10 == 0:
|
|
logger.info(f"Migrated {migrated:,}/{total:,} documents")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Batch insert failed: {e}")
|
|
progress.failed_documents += len(batch)
|
|
# Continue with next batch
|
|
|
|
return migrated
|
|
|
|
def get_schema_preview(self, es_index: str) -> dict[str, Any]:
|
|
"""
|
|
Get a preview of schema analysis without executing migration.
|
|
|
|
Args:
|
|
es_index: Elasticsearch index name
|
|
|
|
Returns:
|
|
Schema analysis information
|
|
"""
|
|
es_mapping = self.es_client.get_index_mapping(es_index)
|
|
analysis = self.schema_converter.analyze_es_mapping(es_mapping)
|
|
column_defs = self.schema_converter.get_column_definitions()
|
|
|
|
return {
|
|
"es_index": es_index,
|
|
"es_mapping": es_mapping,
|
|
"analysis": analysis,
|
|
"ob_columns": column_defs,
|
|
"vector_fields": self.schema_converter.get_vector_fields(),
|
|
"total_columns": len(column_defs),
|
|
}
|
|
|
|
def get_data_preview(
|
|
self,
|
|
es_index: str,
|
|
sample_size: int = 5,
|
|
kb_id: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Get sample documents from ES for preview.
|
|
|
|
Args:
|
|
es_index: ES index name
|
|
sample_size: Number of samples
|
|
kb_id: Optional KB filter
|
|
"""
|
|
query = None
|
|
if kb_id:
|
|
query = {"term": {"kb_id": kb_id}}
|
|
return self.es_client.get_sample_documents(es_index, sample_size, query=query)
|
|
|
|
def list_knowledge_bases(self, es_index: str) -> list[str]:
|
|
"""
|
|
List all knowledge base IDs in an ES index.
|
|
|
|
Args:
|
|
es_index: ES index name
|
|
|
|
Returns:
|
|
List of kb_id values
|
|
"""
|
|
try:
|
|
agg_result = self.es_client.aggregate_field(es_index, "kb_id")
|
|
return [bucket["key"] for bucket in agg_result.get("buckets", [])]
|
|
except Exception as e:
|
|
logger.warning(f"Failed to list knowledge bases: {e}")
|
|
return []
|