zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-17 05:07:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: reformat all code for lefthook using ruff and gofmt (#16585)

Browse Source
This commit is contained in:
Wang Qi
2026-07-03 12:53:39 +08:00
committed by GitHub
parent 19fcb4a981
commit 6a4b9be426
588 changed files with 11123 additions and 15412 deletions
Download Patch File Download Diff File Expand all files Collapse all files

714
tools/scripts/db_schema_sync.py
View File

File diff suppressed because it is too large Load Diff

417
tools/scripts/mysql_migration.py
View File

@@ -49,10 +49,7 @@ PROJECT_BASE = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(_
sys.path.insert(0, PROJECT_BASE)
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s'
)
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
@@ -61,34 +58,34 @@ MIGRATION_DB_VERSION_MARKER = "mysql_migration.database.version"
class MigrationConfig:
"""Configuration for MySQL connection"""
def __init__(self, host: str = 'localhost', port: int = 3306,
user: str = 'root', password: str = '', database: str = 'rag_flow'):
def __init__(self, host: str = "localhost", port: int = 3306, user: str = "root", password: str = "", database: str = "rag_flow"):
self.host = host
self.port = port
self.user = user
self.password = password
self.database = database
@classmethod
def from_config_file(cls, config_path: str) -> 'MigrationConfig':
def from_config_file(cls, config_path: str) -> "MigrationConfig":
"""Load configuration from YAML config file"""
try:
from ruamel.yaml import YAML
yaml = YAML(typ="safe", pure=True)
with open(config_path, 'r') as f:
with open(config_path, "r") as f:
config = yaml.load(f)
# Try to get database config
db_config = config.get('database', config.get('mysql', {}))
db_config = config.get("database", config.get("mysql", {}))
return cls(
host=db_config.get('host', 'localhost'),
port=db_config.get('port', 3306),
user=db_config.get('user', 'root'),
password=db_config.get('password', ''),
database=db_config.get('name', db_config.get('database', 'rag_flow'))
host=db_config.get("host", "localhost"),
port=db_config.get("port", 3306),
user=db_config.get("user", "root"),
password=db_config.get("password", ""),
database=db_config.get("name", db_config.get("database", "rag_flow")),
)
except Exception as e:
logger.warning(f"Failed to load config file: {e}, using defaults")
@@ -97,30 +94,25 @@ class MigrationConfig:
class MigrationStats:
"""Track migration statistics"""
def __init__(self):
self.tables_operated = []
self.rows_processed = 0
self.start_time = None
self.end_time = None
self.stage_stats = []
def start(self):
self.start_time = time.time()
def end(self):
self.end_time = time.time()
def add_stage_stats(self, stage_name: str, tables: list, rows: int, duration: float):
self.stage_stats.append({
'stage': stage_name,
'tables': tables,
'rows': rows,
'duration': duration
})
self.stage_stats.append({"stage": stage_name, "tables": tables, "rows": rows, "duration": duration})
self.tables_operated.extend(tables)
self.rows_processed += rows
def print_summary(self):
duration = self.end_time - self.start_time if self.end_time and self.start_time else 0
logger.info("=" * 60)
@@ -132,52 +124,36 @@ class MigrationStats:
logger.info("-" * 60)
logger.info("Stage Details:")
for stat in self.stage_stats:
logger.info(f" [{stat['stage']}] Tables: {', '.join(stat['tables'])}, "
f"Rows: {stat['rows']}, Duration: {stat['duration']:.2f}s")
logger.info(f" [{stat['stage']}] Tables: {', '.join(stat['tables'])}, Rows: {stat['rows']}, Duration: {stat['duration']:.2f}s")
logger.info("=" * 60)
class MigrationDatabase:
"""Database wrapper for migrations"""
def __init__(self, config: MigrationConfig):
self.config = config
self.db = MySQLDatabase(
config.database,
host=config.host,
port=config.port,
user=config.user,
password=config.password,
charset='utf8mb4'
)
self.db = MySQLDatabase(config.database, host=config.host, port=config.port, user=config.user, password=config.password, charset="utf8mb4")
self.migrator = MySQLMigrator(self.db)
def connect(self):
self.db.connect()
logger.info(f"Connected to MySQL database: {self.config.database}")
def close(self):
if not self.db.is_closed():
self.db.close()
logger.info("Database connection closed")
def execute_sql(self, sql: str, params=None):
return self.db.execute_sql(sql, params)
def table_exists(self, table_name: str) -> bool:
cursor = self.execute_sql(
"SELECT COUNT(*) FROM information_schema.tables "
"WHERE table_schema = %s AND table_name = %s",
(self.config.database, table_name)
)
cursor = self.execute_sql("SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = %s AND table_name = %s", (self.config.database, table_name))
return cursor.fetchone()[0] > 0
def column_exists(self, table_name: str, column_name: str) -> bool:
cursor = self.execute_sql(
"SELECT COUNT(*) FROM information_schema.columns "
"WHERE table_schema = %s AND table_name = %s AND column_name = %s",
(self.config.database, table_name, column_name)
)
cursor = self.execute_sql("SELECT COUNT(*) FROM information_schema.columns WHERE table_schema = %s AND table_name = %s AND column_name = %s", (self.config.database, table_name, column_name))
return cursor.fetchone()[0] > 0
def get_system_setting_value(self, name: str) -> str | None:
@@ -252,17 +228,19 @@ def should_skip_migration(current_db_version: str | None, target_version: str) -
# Define model classes for migration (not importing from api.db.db_models)
class BaseModel(Model):
"""Base model for migration tables"""
create_time = BigIntegerField(null=True, index=True)
create_date = DateTimeField(null=True, index=True)
update_time = BigIntegerField(null=True, index=True)
update_date = DateTimeField(null=True, index=True)
class Meta:
database = None # Will be set dynamically
class TenantLLM(BaseModel):
"""Tenant LLM model (source table)"""
id = PrimaryKeyField()
tenant_id = CharField(max_length=32, null=False, index=True)
llm_factory = CharField(max_length=128, null=False, index=True)
@@ -273,7 +251,7 @@ class TenantLLM(BaseModel):
max_tokens = IntegerField(default=8192, index=True)
used_tokens = IntegerField(default=0, index=True)
status = CharField(max_length=1, null=False, default="1", index=True)
class Meta:
table_name = "tenant_llm"
database = None
@@ -281,10 +259,11 @@ class TenantLLM(BaseModel):
class TenantModelProvider(BaseModel):
"""Tenant Model Provider model (target table)"""
id = CharField(max_length=32, primary_key=True)
provider_name = CharField(max_length=128, null=False, index=True)
tenant_id = CharField(max_length=32, null=False, index=True)
class Meta:
table_name = "tenant_model_provider"
database = None
@@ -292,25 +271,26 @@ class TenantModelProvider(BaseModel):
class MigrationStage:
"""Base class for migration stages"""
name = "base_stage"
description = "Base migration stage"
source_tables = []
target_tables = []
def __init__(self, db: MigrationDatabase, dry_run: bool = True, create_table_only: bool = False):
self.db = db
self.dry_run = dry_run
self.create_table_only = create_table_only
self._noop_completes_migration = False
def check(self) -> bool:
"""Check if migration is needed"""
raise NotImplementedError
def execute(self) -> tuple[int, list]:
"""Execute migration, returns (rows_affected, tables_operated)"""
raise NotImplementedError
def create_target_table(self):
"""Create target table (override in subclass if needed)"""
pass
@@ -324,72 +304,66 @@ class MigrationStage:
class TenantModelProviderStage(MigrationStage):
"""Migrate tenant_llm to tenant_model_provider"""
name = "tenant_model_provider"
description = "Migrate tenant_llm.llm_factory to tenant_model_provider.provider_name"
source_tables = ["tenant_llm"]
target_tables = ["tenant_model_provider"]
def current_timestamp(self) -> int:
return int(time.time())
def generate_uuid(self) -> str:
"""Generate 32-character UUID1"""
return uuid.uuid1().hex
def check(self) -> bool:
"""Check if migration is needed"""
# Check if source table exists
if not self.db.table_exists("tenant_llm"):
logger.warning("Source table 'tenant_llm' does not exist")
return False
# Check if target table exists
if not self.db.table_exists("tenant_model_provider"):
if self.dry_run:
logger.info("[DRY RUN] Target table 'tenant_model_provider' does not exist. "
"Use --execute to create and populate the table.")
logger.info("[DRY RUN] Target table 'tenant_model_provider' does not exist. Use --execute to create and populate the table.")
return False
logger.info("Target table 'tenant_model_provider' does not exist, will create")
return True
# Check if there's data to migrate
cursor = self.db.execute_sql(
"SELECT COUNT(*) FROM tenant_llm t1 "
"WHERE NOT EXISTS ("
" SELECT 1 FROM tenant_model_provider t2 "
" WHERE t2.tenant_id = t1.tenant_id AND t2.provider_name = t1.llm_factory"
")"
"SELECT COUNT(*) FROM tenant_llm t1 WHERE NOT EXISTS ( SELECT 1 FROM tenant_model_provider t2 WHERE t2.tenant_id = t1.tenant_id AND t2.provider_name = t1.llm_factory)"
)
count = cursor.fetchone()[0]
if count == 0:
self.mark_noop_completes_migration()
logger.info("No new data to migrate from tenant_llm to tenant_model_provider")
return False
logger.info(f"Found {count} rows to migrate from tenant_llm to tenant_model_provider")
return True
def execute(self) -> tuple[int, list]:
"""Execute migration"""
current_ts = self.current_timestamp()
rows_inserted = 0
# Check if target table exists
if not self.db.table_exists("tenant_model_provider"):
if self.dry_run:
logger.info("[DRY RUN] Target table 'tenant_model_provider' does not exist. "
"Use --execute to create and populate the table.")
logger.info("[DRY RUN] Target table 'tenant_model_provider' does not exist. Use --execute to create and populate the table.")
return 0, []
logger.info("Target table 'tenant_model_provider' does not exist, will create")
self.create_target_table()
# If create_table_only mode, skip data migration
if self.create_table_only:
logger.info("[CREATE TABLE ONLY] Target table created/verified, skipping data migration")
return 0, self.target_tables
# Get distinct tenant_id, llm_factory pairs that don't exist in target
cursor = self.db.execute_sql(
"SELECT DISTINCT tenant_id, llm_factory FROM tenant_llm t1 "
@@ -398,48 +372,50 @@ class TenantModelProviderStage(MigrationStage):
" WHERE t2.tenant_id = t1.tenant_id AND t2.provider_name = t1.llm_factory"
")"
)
records = cursor.fetchall()
if not records:
logger.info("No records to migrate")
return 0, []
logger.info(f"Migrating {len(records)} unique tenant_id/llm_factory pairs...")
if self.dry_run:
logger.info(f"[DRY RUN] Would insert {len(records)} records")
return len(records), self.target_tables
# Insert records in batches with parameterized SQL to avoid quote breakage/injection
batch_size = 100
for i in range(0, len(records), batch_size):
batch = records[i:i + batch_size]
batch = records[i : i + batch_size]
placeholders = []
params = []
for tenant_id, llm_factory in batch:
record_id = self.generate_uuid()
placeholders.append("(%s, %s, %s, %s, FROM_UNIXTIME(%s), %s, FROM_UNIXTIME(%s))")
params.extend([
record_id,
llm_factory,
tenant_id,
current_ts * 1000,
current_ts,
current_ts * 1000,
current_ts,
])
params.extend(
[
record_id,
llm_factory,
tenant_id,
current_ts * 1000,
current_ts,
current_ts * 1000,
current_ts,
]
)
insert_sql = f"""
INSERT INTO tenant_model_provider
INSERT INTO tenant_model_provider
(id, provider_name, tenant_id, create_time, create_date, update_time, update_date)
VALUES {', '.join(placeholders)}
VALUES {", ".join(placeholders)}
"""
self.db.execute_sql(insert_sql, params)
rows_inserted += len(batch)
logger.info(f"Inserted batch {i // batch_size + 1}: {len(batch)} records")
return rows_inserted, self.target_tables
def create_target_table(self):
"""Create tenant_model_provider table"""
create_sql = """
@@ -485,18 +461,15 @@ class TenantModelInstanceStage(MigrationStage):
# Check if tenant_model_provider exists (dependency)
if not self.db.table_exists("tenant_model_provider"):
if self.dry_run:
logger.info("[DRY RUN] Dependency table 'tenant_model_provider' does not exist. "
"Run 'tenant_model_provider' stage first or use --execute.")
logger.info("[DRY RUN] Dependency table 'tenant_model_provider' does not exist. Run 'tenant_model_provider' stage first or use --execute.")
return False
logger.warning("Dependency table 'tenant_model_provider' does not exist. "
"Please run 'tenant_model_provider' stage first.")
logger.warning("Dependency table 'tenant_model_provider' does not exist. Please run 'tenant_model_provider' stage first.")
return False
# Check if target table exists
if not self.db.table_exists("tenant_model_instance"):
if self.dry_run:
logger.info("[DRY RUN] Target table 'tenant_model_instance' does not exist. "
"Use --execute to create and populate the table.")
logger.info("[DRY RUN] Target table 'tenant_model_instance' does not exist. Use --execute to create and populate the table.")
return False
logger.info("Target table 'tenant_model_instance' does not exist, will create")
return True
@@ -531,15 +504,13 @@ class TenantModelInstanceStage(MigrationStage):
# Check if tenant_model_provider exists (dependency)
if not self.db.table_exists("tenant_model_provider"):
logger.error("Dependency table 'tenant_model_provider' does not exist. "
"Please run 'tenant_model_provider' stage first.")
logger.error("Dependency table 'tenant_model_provider' does not exist. Please run 'tenant_model_provider' stage first.")
return 0, []
# Check if target table exists
if not self.db.table_exists("tenant_model_instance"):
if self.dry_run:
logger.info("[DRY RUN] Target table 'tenant_model_instance' does not exist. "
"Use --execute to create and populate the table.")
logger.info("[DRY RUN] Target table 'tenant_model_instance' does not exist. Use --execute to create and populate the table.")
return 0, []
logger.info("Target table 'tenant_model_instance' does not exist, will create")
self.create_target_table()
@@ -588,22 +559,24 @@ class TenantModelInstanceStage(MigrationStage):
# Insert records in batches
batch_size = 100
for i in range(0, len(records), batch_size):
batch = records[i:i + batch_size]
batch = records[i : i + batch_size]
values = []
for tenant_id, llm_factory, api_key, status, provider_id in batch:
record_id = self.generate_uuid()
instance_name = "default"
api_key_escaped = api_key.replace("'", "''") if api_key else ""
status_val = "active" if status in ["1", "active", "enable"] else "inactive"
values.append(f"('{record_id}', '{instance_name}', '{provider_id}', "
f"'{api_key_escaped}', '{status_val}', "
f"{current_ts * 1000}, FROM_UNIXTIME({current_ts}), "
f"{current_ts * 1000}, FROM_UNIXTIME({current_ts}))")
values.append(
f"('{record_id}', '{instance_name}', '{provider_id}', "
f"'{api_key_escaped}', '{status_val}', "
f"{current_ts * 1000}, FROM_UNIXTIME({current_ts}), "
f"{current_ts * 1000}, FROM_UNIXTIME({current_ts}))"
)
insert_sql = f"""
INSERT INTO tenant_model_instance
INSERT INTO tenant_model_instance
(id, instance_name, provider_id, api_key, status, create_time, create_date, update_time, update_date)
VALUES {', '.join(values)}
VALUES {", ".join(values)}
"""
self.db.execute_sql(insert_sql)
rows_inserted += len(batch)
@@ -688,11 +661,7 @@ class TenantModelInstanceStage(MigrationStage):
seen[canonical] = rec
else:
dup_count += 1
logger.debug(
f"Dedup api_key for tenant={tenant_id}, factory={llm_factory}, "
f"provider={provider_id}: keeping '{api_key[:20]}...', "
f"dropping '{seen[canonical][2][:20]}...'"
)
logger.debug(f"Dedup api_key for tenant={tenant_id}, factory={llm_factory}, provider={provider_id}: keeping '{api_key[:20]}...', dropping '{seen[canonical][2][:20]}...'")
deduped.extend(seen.values())
if dup_count > 0:
@@ -770,28 +739,23 @@ class TenantModelStage(MigrationStage):
# Check if tenant_model_provider exists (dependency)
if not self.db.table_exists("tenant_model_provider"):
if self.dry_run:
logger.info("[DRY RUN] Dependency table 'tenant_model_provider' does not exist. "
"Run 'tenant_model_provider' stage first or use --execute.")
logger.info("[DRY RUN] Dependency table 'tenant_model_provider' does not exist. Run 'tenant_model_provider' stage first or use --execute.")
return False
logger.warning("Dependency table 'tenant_model_provider' does not exist. "
"Please run 'tenant_model_provider' stage first.")
logger.warning("Dependency table 'tenant_model_provider' does not exist. Please run 'tenant_model_provider' stage first.")
return False
# Check if tenant_model_instance exists (dependency)
if not self.db.table_exists("tenant_model_instance"):
if self.dry_run:
logger.info("[DRY RUN] Dependency table 'tenant_model_instance' does not exist. "
"Run 'tenant_model_instance' stage first or use --execute.")
logger.info("[DRY RUN] Dependency table 'tenant_model_instance' does not exist. Run 'tenant_model_instance' stage first or use --execute.")
return False
logger.warning("Dependency table 'tenant_model_instance' does not exist. "
"Please run 'tenant_model_instance' stage first.")
logger.warning("Dependency table 'tenant_model_instance' does not exist. Please run 'tenant_model_instance' stage first.")
return False
# Check if target table exists
if not self.db.table_exists("tenant_model"):
if self.dry_run:
logger.info("[DRY RUN] Target table 'tenant_model' does not exist. "
"Use --execute to create and populate the table.")
logger.info("[DRY RUN] Target table 'tenant_model' does not exist. Use --execute to create and populate the table.")
return False
logger.info("Target table 'tenant_model' does not exist, will create")
return True
@@ -827,21 +791,18 @@ class TenantModelStage(MigrationStage):
# Check if tenant_model_provider exists (dependency)
if not self.db.table_exists("tenant_model_provider"):
logger.error("Dependency table 'tenant_model_provider' does not exist. "
"Please run 'tenant_model_provider' stage first.")
logger.error("Dependency table 'tenant_model_provider' does not exist. Please run 'tenant_model_provider' stage first.")
return 0, []
# Check if tenant_model_instance exists (dependency)
if not self.db.table_exists("tenant_model_instance"):
logger.error("Dependency table 'tenant_model_instance' does not exist. "
"Please run 'tenant_model_instance' stage first.")
logger.error("Dependency table 'tenant_model_instance' does not exist. Please run 'tenant_model_instance' stage first.")
return 0, []
# Check if target table exists
if not self.db.table_exists("tenant_model"):
if self.dry_run:
logger.info("[DRY RUN] Target table 'tenant_model' does not exist. "
"Use --execute to create and populate the table.")
logger.info("[DRY RUN] Target table 'tenant_model' does not exist. Use --execute to create and populate the table.")
return 0, []
logger.info("Target table 'tenant_model' does not exist, will create")
self.create_target_table()
@@ -891,8 +852,7 @@ class TenantModelStage(MigrationStage):
if self.dry_run:
logger.info(f"[DRY RUN] Would insert {len(resolved_records)} records")
for source_id, llm_name, provider_id, instance_id, model_type, status, api_key in resolved_records[:5]:
logger.info(f" model_name={llm_name}, provider_id={provider_id}, "
f"instance_id={instance_id}, model_type={model_type}")
logger.info(f" model_name={llm_name}, provider_id={provider_id}, instance_id={instance_id}, model_type={model_type}")
if len(resolved_records) > 5:
logger.info(f" ... and {len(resolved_records) - 5} more records")
return len(resolved_records), self.target_tables
@@ -900,7 +860,7 @@ class TenantModelStage(MigrationStage):
# Insert records in batches
batch_size = 100
for i in range(0, len(resolved_records), batch_size):
batch = resolved_records[i:i + batch_size]
batch = resolved_records[i : i + batch_size]
values = []
for source_id, llm_name, provider_id, instance_id, model_type, status, api_key in batch:
record_id = self.generate_uuid()
@@ -910,17 +870,19 @@ class TenantModelStage(MigrationStage):
# Extract is_tools from api_key JSON and put it in extra
extra = self._extract_extra_from_api_key(api_key)
extra_escaped = extra.replace("'", "''") if extra else "{}"
values.append(f"('{record_id}', '{model_name_escaped}', '{provider_id}', "
f"'{instance_id}', '{model_type_escaped}', '{status_val}', "
f"'{extra_escaped}', "
f"{current_ts * 1000}, FROM_UNIXTIME({current_ts}), "
f"{current_ts * 1000}, FROM_UNIXTIME({current_ts}))")
values.append(
f"('{record_id}', '{model_name_escaped}', '{provider_id}', "
f"'{instance_id}', '{model_type_escaped}', '{status_val}', "
f"'{extra_escaped}', "
f"{current_ts * 1000}, FROM_UNIXTIME({current_ts}), "
f"{current_ts * 1000}, FROM_UNIXTIME({current_ts}))"
)
insert_sql = f"""
INSERT INTO tenant_model
INSERT INTO tenant_model
(id, model_name, provider_id, instance_id, model_type, status, extra,
create_time, create_date, update_time, update_date)
VALUES {', '.join(values)}
VALUES {", ".join(values)}
"""
self.db.execute_sql(insert_sql)
rows_inserted += len(batch)
@@ -937,9 +899,7 @@ class TenantModelStage(MigrationStage):
Returns:
dict mapping (provider_id, canonical_api_key) -> instance_id
"""
cursor = self.db.execute_sql(
"SELECT id, provider_id, api_key FROM tenant_model_instance"
)
cursor = self.db.execute_sql("SELECT id, provider_id, api_key FROM tenant_model_instance")
lookup = {}
for instance_id, provider_id, api_key in cursor.fetchall():
canonical = TenantModelInstanceStage._strip_is_tools_from_api_key(api_key)
@@ -974,7 +934,9 @@ class TenantModelStage(MigrationStage):
# entropy to be useful to an attacker who reads the log.
logger.warning(
"No matching instance for tenant_llm id=%s provider_id=%s llm_name=%s",
source_id, provider_id, llm_name,
source_id,
provider_id,
llm_name,
)
if skipped > 0:
@@ -1102,9 +1064,7 @@ class ModelIdConfigStage(MigrationStage):
normalized = {}
for key, item in value.items():
key_path = path + (str(key),)
should_normalize = key in self.model_id_fields or (
key in self.search_config_model_id_fields and "search_config" in path
)
should_normalize = key in self.model_id_fields or (key in self.search_config_model_id_fields and "search_config" in path)
if should_normalize:
normalized_item, item_changed = self.normalize_model_id(item)
else:
@@ -1154,8 +1114,7 @@ class ModelIdConfigStage(MigrationStage):
def iter_string_changes(self):
for table_name, column_name in self.existing_columns(self.string_columns):
cursor = self.db.execute_sql(
f"SELECT id, `{column_name}` FROM `{table_name}` "
f"WHERE `{column_name}` IS NOT NULL AND `{column_name}` != '' AND `{column_name}` LIKE %s",
f"SELECT id, `{column_name}` FROM `{table_name}` WHERE `{column_name}` IS NOT NULL AND `{column_name}` != '' AND `{column_name}` LIKE %s",
("%@%",),
)
while True:
@@ -1170,8 +1129,7 @@ class ModelIdConfigStage(MigrationStage):
def iter_json_changes(self):
for table_name, column_name in self.existing_columns(self.json_columns):
cursor = self.db.execute_sql(
f"SELECT id, `{column_name}` FROM `{table_name}` "
f"WHERE `{column_name}` IS NOT NULL AND `{column_name}` != '' AND `{column_name}` LIKE %s",
f"SELECT id, `{column_name}` FROM `{table_name}` WHERE `{column_name}` IS NOT NULL AND `{column_name}` != '' AND `{column_name}` LIKE %s",
("%@%",),
)
while True:
@@ -1271,10 +1229,10 @@ class ModelIdConfigStage(MigrationStage):
# Registry of available migration stages
MIGRATION_STAGES = {
'tenant_model_provider': TenantModelProviderStage,
'tenant_model_instance': TenantModelInstanceStage,
'tenant_model': TenantModelStage,
'model_id_config': ModelIdConfigStage,
"tenant_model_provider": TenantModelProviderStage,
"tenant_model_instance": TenantModelInstanceStage,
"tenant_model": TenantModelStage,
"model_id_config": ModelIdConfigStage,
}
@@ -1298,9 +1256,9 @@ def run_migration(
"""Run migration with specified stages"""
stats = MigrationStats()
stats.start()
db = MigrationDatabase(config)
try:
db.connect()
@@ -1325,26 +1283,26 @@ def run_migration(
"Database migration version marker is not set, target version is %s",
database_version,
)
total_stages = len(stages)
all_stages_completed = True
for idx, stage_name in enumerate(stages, 1):
logger.info(f"{'=' * 60}")
logger.info(f"Stage [{idx}/{total_stages}]: {stage_name}")
logger.info(f"{'=' * 60}")
if stage_name not in MIGRATION_STAGES:
logger.error(f"Unknown stage: {stage_name}")
stats.add_stage_stats(stage_name, [], 0, 0)
all_stages_completed = False
continue
stage_cls = MIGRATION_STAGES[stage_name]
stage = stage_cls(db, dry_run=dry_run, create_table_only=create_table_only)
stage_start = time.time()
# For create_table_only mode, skip check and directly execute
if create_table_only:
logger.info("[CREATE TABLE ONLY] Skipping check, will create/verify target table")
@@ -1355,22 +1313,16 @@ def run_migration(
logger.info(f"Stage '{stage_name}' check: no migration needed")
stats.add_stage_stats(stage_name, [], 0, time.time() - stage_start)
continue
# Execute migration
rows, tables = stage.execute()
stage_duration = time.time() - stage_start
stats.add_stage_stats(stage_name, tables, rows, stage_duration)
logger.info(f"Stage '{stage_name}' completed: {rows} rows in {stage_duration:.2f}s")
if (
mark_database_version_on_success
and not dry_run
and not create_table_only
and database_version
and all_stages_completed
):
if mark_database_version_on_success and not dry_run and not create_table_only and database_version and all_stages_completed:
db.set_database_version(database_version)
logger.info("Marked database migration version as %s", database_version)
@@ -1421,7 +1373,7 @@ def mark_database_version(config: MigrationConfig, version: str) -> None:
def main():
parser = argparse.ArgumentParser(
description='MySQL Data Migration Tool',
description="MySQL Data Migration Tool",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
@@ -1433,16 +1385,16 @@ Examples:
# Mark database version separately
python mysql_migration.py --mark-database-version --database-version v0.26.3 --config /path/to/config.yaml
# Dry run (default - check only, no write) with config file
python mysql_migration.py --stages tenant_model_provider --config /path/to/config.yaml
# Dry run with command line MySQL connection
python mysql_migration.py --stages tenant_model_provider --host localhost --port 3306 --user root --password secret
# Create target tables only (no data migration)
python mysql_migration.py --stages tenant_model_provider --config /path/to/config.yaml --create-table-only
# Execute full migration (create tables and migrate data)
python mysql_migration.py --stages tenant_model_provider --config /path/to/config.yaml --execute
@@ -1451,79 +1403,61 @@ Examples:
# Execute migration and mark the database version when all stages succeed
python mysql_migration.py --stages tenant_model_provider,tenant_model_instance,tenant_model,model_id_config --config /path/to/config.yaml --execute --database-version v0.26.3 --mark-database-version-on-success
# Normalize legacy model IDs in stored configs
python mysql_migration.py --stages model_id_config --config /path/to/config.yaml --execute
# Run multiple stages
python mysql_migration.py --stages stage1,stage2,stage3 --config /path/to/config.yaml --execute
"""
""",
)
# MySQL connection options
parser.add_argument('--host', type=str, default='localhost',
help='MySQL host (default: localhost)')
parser.add_argument('--port', type=int, default=3306,
help='MySQL port (default: 3306)')
parser.add_argument('--user', type=str, default='root',
help='MySQL user (default: root)')
parser.add_argument('--password', type=str, default='',
help='MySQL password (default: empty)')
parser.add_argument('--database', type=str, default='rag_flow',
help='MySQL database name (default: rag_flow)')
parser.add_argument("--host", type=str, default="localhost", help="MySQL host (default: localhost)")
parser.add_argument("--port", type=int, default=3306, help="MySQL port (default: 3306)")
parser.add_argument("--user", type=str, default="root", help="MySQL user (default: root)")
parser.add_argument("--password", type=str, default="", help="MySQL password (default: empty)")
parser.add_argument("--database", type=str, default="rag_flow", help="MySQL database name (default: rag_flow)")
# Configuration options
parser.add_argument('--config', '-c', type=str, help='Path to YAML config file')
parser.add_argument("--config", "-c", type=str, help="Path to YAML config file")
# Migration options
parser.add_argument('--stages', '-s', type=str, help='Comma-separated list of stages to run')
parser.add_argument('--list-stages', '-l', action='store_true', help='List available stages')
parser.add_argument('--check-database-version', action='store_true',
help='Check whether migration is needed for the target database version')
parser.add_argument('--mark-database-version', action='store_true',
help='Write the database migration version marker and exit')
parser.add_argument('--database-version', type=str, metavar='VERSION',
help='Database migration version used by check/mark commands and as the migration threshold for --stages')
parser.add_argument('--mark-database-version-on-success', action='store_true',
help='When used with --stages and --execute, write --database-version after all stages succeed')
parser.add_argument('--execute', '-e', action='store_true', default=False,
help='Execute full migration: create tables and migrate data')
parser.add_argument('--create-table-only', action='store_true', default=False,
help='Only create target tables, skip data migration')
parser.add_argument("--stages", "-s", type=str, help="Comma-separated list of stages to run")
parser.add_argument("--list-stages", "-l", action="store_true", help="List available stages")
parser.add_argument("--check-database-version", action="store_true", help="Check whether migration is needed for the target database version")
parser.add_argument("--mark-database-version", action="store_true", help="Write the database migration version marker and exit")
parser.add_argument("--database-version", type=str, metavar="VERSION", help="Database migration version used by check/mark commands and as the migration threshold for --stages")
parser.add_argument("--mark-database-version-on-success", action="store_true", help="When used with --stages and --execute, write --database-version after all stages succeed")
parser.add_argument("--execute", "-e", action="store_true", default=False, help="Execute full migration: create tables and migrate data")
parser.add_argument("--create-table-only", action="store_true", default=False, help="Only create target tables, skip data migration")
args = parser.parse_args()
# List stages and exit
if args.list_stages:
list_available_stages()
return
# Load configuration: command line args take precedence over config file
if args.config:
config = MigrationConfig.from_config_file(args.config)
# Override with command line args if provided
if args.host != 'localhost':
if args.host != "localhost":
config.host = args.host
if args.port != 3306:
config.port = args.port
if args.user != 'root':
if args.user != "root":
config.user = args.user
if args.password != '':
if args.password != "":
config.password = args.password
if args.database != 'rag_flow':
if args.database != "rag_flow":
config.database = args.database
else:
# Use command line args directly
config = MigrationConfig(
host=args.host,
port=args.port,
user=args.user,
password=args.password,
database=args.database
)
logger.info(f"MySQL Configuration: host={config.host}, port={config.port}, "
f"user={config.user}, database={config.database}")
config = MigrationConfig(host=args.host, port=args.port, user=args.user, password=args.password, database=args.database)
logger.info(f"MySQL Configuration: host={config.host}, port={config.port}, user={config.user}, database={config.database}")
if args.check_database_version and args.mark_database_version:
logger.error("--check-database-version and --mark-database-version are mutually exclusive")
@@ -1545,7 +1479,7 @@ Examples:
if args.mark_database_version_on_success and not args.database_version:
logger.error("--mark-database-version-on-success requires --database-version")
sys.exit(1)
# Three mutually exclusive modes: dry-run (default), create-table-only, execute
if args.execute and args.create_table_only:
logger.error("--execute and --create-table-only are mutually exclusive")
@@ -1555,10 +1489,10 @@ Examples:
logger.error("No stages specified. Use --stages to specify stages or --list-stages to see available stages.")
sys.exit(1)
stages = [s.strip() for s in args.stages.split(',')]
stages = [s.strip() for s in args.stages.split(",")]
dry_run = True
create_table_only = False
if args.create_table_only:
logger.info("Running in CREATE TABLE ONLY mode (create tables, no data migration)")
dry_run = False
@@ -1567,9 +1501,8 @@ Examples:
logger.info("Running in EXECUTE mode (create tables and migrate data)")
dry_run = False
else:
logger.info("Running in DRY-RUN mode (check only, no write). "
"Use --create-table-only to create tables, or --execute for full migration.")
logger.info("Running in DRY-RUN mode (check only, no write). Use --create-table-only to create tables, or --execute for full migration.")
run_migration(
config=config,
stages=stages,
@@ -1580,5 +1513,5 @@ Examples:
)
if __name__ == '__main__':
if __name__ == "__main__":
main()