mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-30 20:49:21 +08:00
1b587013d8
### What problem does this PR solve? - Remove unused imports (Mock, patch, MagicMock, json, os, RAGFLOW_COLUMNS, VECTOR_FIELD_PATTERN) from multiple files - Replace f-string formatting with regular strings for console output messages in cli.py - Clean up unnecessary imports that were no longer being used in the codebase ### Type of change - [x] Refactoring
350 lines
11 KiB
Python
350 lines
11 KiB
Python
"""
|
|
Data verification for RAGFlow migration.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from .es_client import ESClient
|
|
from .ob_client import OBClient
|
|
from .schema import ARRAY_COLUMNS, JSON_COLUMNS
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class VerificationResult:
|
|
"""Migration verification result."""
|
|
|
|
es_index: str
|
|
ob_table: str
|
|
|
|
# Counts
|
|
es_count: int = 0
|
|
ob_count: int = 0
|
|
count_match: bool = False
|
|
count_diff: int = 0
|
|
|
|
# Sample verification
|
|
sample_size: int = 0
|
|
samples_verified: int = 0
|
|
samples_matched: int = 0
|
|
sample_match_rate: float = 0.0
|
|
|
|
# Mismatches
|
|
missing_in_ob: list[str] = field(default_factory=list)
|
|
data_mismatches: list[dict[str, Any]] = field(default_factory=list)
|
|
|
|
# Overall
|
|
passed: bool = False
|
|
message: str = ""
|
|
|
|
|
|
class MigrationVerifier:
|
|
"""Verify RAGFlow migration data consistency."""
|
|
|
|
# Fields to compare for verification
|
|
VERIFY_FIELDS = [
|
|
"id", "kb_id", "doc_id", "docnm_kwd", "content_with_weight",
|
|
"available_int", "create_time",
|
|
]
|
|
|
|
def __init__(
|
|
self,
|
|
es_client: ESClient,
|
|
ob_client: OBClient,
|
|
):
|
|
"""
|
|
Initialize verifier.
|
|
|
|
Args:
|
|
es_client: Elasticsearch client
|
|
ob_client: OceanBase client
|
|
"""
|
|
self.es_client = es_client
|
|
self.ob_client = ob_client
|
|
|
|
def verify(
|
|
self,
|
|
es_index: str,
|
|
ob_table: str,
|
|
sample_size: int = 100,
|
|
primary_key: str = "id",
|
|
verify_fields: list[str] | None = None,
|
|
) -> VerificationResult:
|
|
"""
|
|
Verify migration by comparing ES and OceanBase data.
|
|
|
|
Args:
|
|
es_index: Elasticsearch index name
|
|
ob_table: OceanBase table name
|
|
sample_size: Number of documents to sample for verification
|
|
primary_key: Primary key column name
|
|
verify_fields: Fields to verify (None = use defaults)
|
|
|
|
Returns:
|
|
VerificationResult with details
|
|
"""
|
|
result = VerificationResult(
|
|
es_index=es_index,
|
|
ob_table=ob_table,
|
|
)
|
|
|
|
if verify_fields is None:
|
|
verify_fields = self.VERIFY_FIELDS
|
|
|
|
# Step 1: Verify document counts
|
|
logger.info("Verifying document counts...")
|
|
|
|
result.es_count = self.es_client.count_documents(es_index)
|
|
result.ob_count = self.ob_client.count_rows(ob_table)
|
|
|
|
result.count_diff = abs(result.es_count - result.ob_count)
|
|
result.count_match = result.count_diff == 0
|
|
|
|
logger.info(
|
|
f"Document counts - ES: {result.es_count}, OB: {result.ob_count}, "
|
|
f"Diff: {result.count_diff}"
|
|
)
|
|
|
|
# Step 2: Sample verification
|
|
result.sample_size = min(sample_size, result.es_count)
|
|
|
|
if result.sample_size > 0:
|
|
logger.info(f"Verifying {result.sample_size} sample documents...")
|
|
self._verify_samples(
|
|
es_index, ob_table, result, primary_key, verify_fields
|
|
)
|
|
|
|
# Step 3: Determine overall result
|
|
self._determine_result(result)
|
|
|
|
logger.info(result.message)
|
|
return result
|
|
|
|
def _verify_samples(
|
|
self,
|
|
es_index: str,
|
|
ob_table: str,
|
|
result: VerificationResult,
|
|
primary_key: str,
|
|
verify_fields: list[str],
|
|
):
|
|
"""Verify sample documents."""
|
|
# Get sample documents from ES
|
|
es_samples = self.es_client.get_sample_documents(
|
|
es_index, result.sample_size
|
|
)
|
|
|
|
for es_doc in es_samples:
|
|
result.samples_verified += 1
|
|
doc_id = es_doc.get("_id") or es_doc.get("id")
|
|
|
|
if not doc_id:
|
|
logger.warning("Document without ID found")
|
|
continue
|
|
|
|
# Get corresponding document from OceanBase
|
|
ob_doc = self.ob_client.get_row_by_id(ob_table, doc_id)
|
|
|
|
if ob_doc is None:
|
|
result.missing_in_ob.append(doc_id)
|
|
continue
|
|
|
|
# Compare documents
|
|
match, differences = self._compare_documents(
|
|
es_doc, ob_doc, verify_fields
|
|
)
|
|
|
|
if match:
|
|
result.samples_matched += 1
|
|
else:
|
|
result.data_mismatches.append({
|
|
"id": doc_id,
|
|
"differences": differences,
|
|
})
|
|
|
|
# Calculate match rate
|
|
if result.samples_verified > 0:
|
|
result.sample_match_rate = result.samples_matched / result.samples_verified
|
|
|
|
def _compare_documents(
|
|
self,
|
|
es_doc: dict[str, Any],
|
|
ob_doc: dict[str, Any],
|
|
verify_fields: list[str],
|
|
) -> tuple[bool, list[dict[str, Any]]]:
|
|
"""
|
|
Compare ES document with OceanBase row.
|
|
|
|
Returns:
|
|
Tuple of (match: bool, differences: list)
|
|
"""
|
|
differences = []
|
|
|
|
for field_name in verify_fields:
|
|
es_value = es_doc.get(field_name)
|
|
ob_value = ob_doc.get(field_name)
|
|
|
|
# Skip if both are None/null
|
|
if es_value is None and ob_value is None:
|
|
continue
|
|
|
|
# Handle special comparisons
|
|
if not self._values_equal(field_name, es_value, ob_value):
|
|
differences.append({
|
|
"field": field_name,
|
|
"es_value": es_value,
|
|
"ob_value": ob_value,
|
|
})
|
|
|
|
return len(differences) == 0, differences
|
|
|
|
def _values_equal(
|
|
self,
|
|
field_name: str,
|
|
es_value: Any,
|
|
ob_value: Any
|
|
) -> bool:
|
|
"""Compare two values with type-aware logic."""
|
|
if es_value is None and ob_value is None:
|
|
return True
|
|
|
|
if es_value is None or ob_value is None:
|
|
# One is None, the other isn't
|
|
# For optional fields, this might be acceptable
|
|
return False
|
|
|
|
# Handle array fields (stored as JSON strings in OB)
|
|
if field_name in ARRAY_COLUMNS:
|
|
if isinstance(ob_value, str):
|
|
try:
|
|
ob_value = json.loads(ob_value)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
if isinstance(es_value, list) and isinstance(ob_value, list):
|
|
return set(str(x) for x in es_value) == set(str(x) for x in ob_value)
|
|
|
|
# Handle JSON fields
|
|
if field_name in JSON_COLUMNS:
|
|
if isinstance(ob_value, str):
|
|
try:
|
|
ob_value = json.loads(ob_value)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
if isinstance(es_value, str):
|
|
try:
|
|
es_value = json.loads(es_value)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
return es_value == ob_value
|
|
|
|
# Handle content_with_weight which might be dict or string
|
|
if field_name == "content_with_weight":
|
|
if isinstance(ob_value, str) and isinstance(es_value, dict):
|
|
try:
|
|
ob_value = json.loads(ob_value)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Handle kb_id which might be list in ES
|
|
if field_name == "kb_id":
|
|
if isinstance(es_value, list) and len(es_value) > 0:
|
|
es_value = es_value[0]
|
|
|
|
# Standard comparison
|
|
return str(es_value) == str(ob_value)
|
|
|
|
def _determine_result(self, result: VerificationResult):
|
|
"""Determine overall verification result."""
|
|
# Allow small count differences (e.g., documents added during migration)
|
|
count_tolerance = 0.01 # 1% tolerance
|
|
count_ok = (
|
|
result.count_match or
|
|
(result.es_count > 0 and result.count_diff / result.es_count <= count_tolerance)
|
|
)
|
|
|
|
if count_ok and result.sample_match_rate >= 0.99:
|
|
result.passed = True
|
|
result.message = (
|
|
f"Verification PASSED. "
|
|
f"ES: {result.es_count:,}, OB: {result.ob_count:,}. "
|
|
f"Sample match rate: {result.sample_match_rate:.2%}"
|
|
)
|
|
elif count_ok and result.sample_match_rate >= 0.95:
|
|
result.passed = True
|
|
result.message = (
|
|
f"Verification PASSED with warnings. "
|
|
f"ES: {result.es_count:,}, OB: {result.ob_count:,}. "
|
|
f"Sample match rate: {result.sample_match_rate:.2%}"
|
|
)
|
|
else:
|
|
result.passed = False
|
|
issues = []
|
|
if not count_ok:
|
|
issues.append(
|
|
f"Count mismatch (ES: {result.es_count}, OB: {result.ob_count}, diff: {result.count_diff})"
|
|
)
|
|
if result.sample_match_rate < 0.95:
|
|
issues.append(f"Low sample match rate: {result.sample_match_rate:.2%}")
|
|
if result.missing_in_ob:
|
|
issues.append(f"{len(result.missing_in_ob)} documents missing in OB")
|
|
result.message = f"Verification FAILED: {'; '.join(issues)}"
|
|
|
|
def generate_report(self, result: VerificationResult) -> str:
|
|
"""Generate a verification report."""
|
|
lines = [
|
|
"",
|
|
"=" * 60,
|
|
"Migration Verification Report",
|
|
"=" * 60,
|
|
f"ES Index: {result.es_index}",
|
|
f"OB Table: {result.ob_table}",
|
|
]
|
|
|
|
lines.extend([
|
|
"",
|
|
"Document Counts:",
|
|
f" Elasticsearch: {result.es_count:,}",
|
|
f" OceanBase: {result.ob_count:,}",
|
|
f" Difference: {result.count_diff:,}",
|
|
f" Match: {'Yes' if result.count_match else 'No'}",
|
|
"",
|
|
"Sample Verification:",
|
|
f" Sample Size: {result.sample_size}",
|
|
f" Verified: {result.samples_verified}",
|
|
f" Matched: {result.samples_matched}",
|
|
f" Match Rate: {result.sample_match_rate:.2%}",
|
|
"",
|
|
])
|
|
|
|
if result.missing_in_ob:
|
|
lines.append(f"Missing in OceanBase ({len(result.missing_in_ob)}):")
|
|
for doc_id in result.missing_in_ob[:5]:
|
|
lines.append(f" - {doc_id}")
|
|
if len(result.missing_in_ob) > 5:
|
|
lines.append(f" ... and {len(result.missing_in_ob) - 5} more")
|
|
lines.append("")
|
|
|
|
if result.data_mismatches:
|
|
lines.append(f"Data Mismatches ({len(result.data_mismatches)}):")
|
|
for mismatch in result.data_mismatches[:3]:
|
|
lines.append(f" - ID: {mismatch['id']}")
|
|
for diff in mismatch.get("differences", [])[:2]:
|
|
lines.append(f" {diff['field']}: ES={diff['es_value']}, OB={diff['ob_value']}")
|
|
if len(result.data_mismatches) > 3:
|
|
lines.append(f" ... and {len(result.data_mismatches) - 3} more")
|
|
lines.append("")
|
|
|
|
lines.extend([
|
|
"=" * 60,
|
|
f"Result: {'PASSED' if result.passed else 'FAILED'}",
|
|
result.message,
|
|
"=" * 60,
|
|
"",
|
|
])
|
|
|
|
return "\n".join(lines)
|