mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-15 05:04:27 +08:00
Feat: migrate script (#13976)
### What problem does this PR solve? Add stage for migrate tenant_llm data into table tenant_model_instance and tenant_model. ### Type of change - [x] Other (please describe): tool script <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Added two new migration stages to move tenant model and instance records into new target tables, with dry-run, full-execute, and "create table only" modes; migration skips already-migrated rows to avoid duplicates. * **Bug Fixes** * Cleaned up migration header logging for clearer output. * **Documentation** * Added usage guide describing stages, options, modes, config format, examples, and expected logs. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
Lynn
@@ -365,9 +365,362 @@ class TenantModelProviderStage(MigrationStage):
|
||||
logger.info("Created tenant_model_provider table")
|
||||
|
||||
|
||||
class TenantModelInstanceStage(MigrationStage):
|
||||
"""Migrate tenant_llm to tenant_model_instance"""
|
||||
|
||||
name = "tenant_model_instance"
|
||||
description = "Migrate tenant_llm to tenant_model_instance with provider_id lookup"
|
||||
source_tables = ["tenant_llm", "tenant_model_provider"]
|
||||
target_tables = ["tenant_model_instance"]
|
||||
|
||||
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 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.")
|
||||
return False
|
||||
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.")
|
||||
return False
|
||||
logger.info("Target table 'tenant_model_instance' does not exist, will create")
|
||||
return True
|
||||
|
||||
# Check if there's data to migrate (distinct by tenant_id, llm_factory, api_key)
|
||||
cursor = self.db.execute_sql(
|
||||
"SELECT COUNT(*) FROM ("
|
||||
" SELECT tl.tenant_id, tl.llm_factory, tl.api_key, tmp.id as provider_id "
|
||||
" FROM tenant_llm tl "
|
||||
" INNER JOIN tenant_model_provider tmp ON tmp.tenant_id = tl.tenant_id AND tmp.provider_name = tl.llm_factory "
|
||||
" WHERE NOT EXISTS ("
|
||||
" SELECT 1 FROM tenant_model_instance tmi "
|
||||
" WHERE tmi.provider_id = tmp.id AND tmi.api_key = tl.api_key"
|
||||
" ) "
|
||||
" GROUP BY tl.tenant_id, tl.llm_factory, tl.api_key, tmp.id"
|
||||
") AS distinct_records"
|
||||
)
|
||||
count = cursor.fetchone()[0]
|
||||
|
||||
if count == 0:
|
||||
logger.info("No new data to migrate from tenant_llm to tenant_model_instance")
|
||||
return False
|
||||
|
||||
logger.info(f"Found {count} rows to migrate from tenant_llm to tenant_model_instance")
|
||||
return True
|
||||
|
||||
def execute(self) -> tuple[int, list]:
|
||||
"""Execute migration"""
|
||||
current_ts = self.current_timestamp()
|
||||
rows_inserted = 0
|
||||
|
||||
# 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.")
|
||||
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.")
|
||||
return 0, []
|
||||
logger.info("Target table 'tenant_model_instance' 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 records from tenant_llm with provider_id lookup
|
||||
# Group by tenant_id, llm_factory, api_key to get distinct records
|
||||
# instance_name = llm_factory, provider_id from tenant_model_provider, api_key from tenant_llm
|
||||
cursor = self.db.execute_sql(
|
||||
"SELECT tl.tenant_id, tl.llm_factory, tl.api_key, MAX(tl.status) as status, tmp.id as provider_id "
|
||||
"FROM tenant_llm tl "
|
||||
"INNER JOIN tenant_model_provider tmp ON tmp.tenant_id = tl.tenant_id AND tmp.provider_name = tl.llm_factory "
|
||||
"WHERE NOT EXISTS ("
|
||||
" SELECT 1 FROM tenant_model_instance tmi "
|
||||
" WHERE tmi.provider_id = tmp.id AND tmi.api_key = tl.api_key"
|
||||
") "
|
||||
"GROUP BY tl.tenant_id, tl.llm_factory, tl.api_key, tmp.id"
|
||||
)
|
||||
|
||||
records = cursor.fetchall()
|
||||
|
||||
if not records:
|
||||
logger.info("No records to migrate")
|
||||
return 0, []
|
||||
|
||||
logger.info(f"Migrating {len(records)} tenant_model_instance records...")
|
||||
|
||||
if self.dry_run:
|
||||
logger.info(f"[DRY RUN] Would insert {len(records)} records")
|
||||
for tenant_id, llm_factory, api_key, status, provider_id in records[:5]:
|
||||
logger.info(f" instance_name={llm_factory}, provider_id={provider_id}, api_key=***")
|
||||
if len(records) > 5:
|
||||
logger.info(f" ... and {len(records) - 5} more records")
|
||||
return len(records), self.target_tables
|
||||
|
||||
# Insert records in batches
|
||||
batch_size = 100
|
||||
for i in range(0, len(records), 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 = llm_factory.replace("'", "''") if llm_factory else ""
|
||||
api_key_escaped = api_key.replace("'", "''") if api_key else ""
|
||||
status_val = status if status else "active"
|
||||
values.append(f"('{record_id}', '{instance_name}', '{provider_id}', "
|
||||
f"'{api_key_escaped}', '{status_val}', "
|
||||
f"{current_ts}, FROM_UNIXTIME({current_ts}), "
|
||||
f"{current_ts}, FROM_UNIXTIME({current_ts}))")
|
||||
|
||||
insert_sql = f"""
|
||||
INSERT INTO tenant_model_instance
|
||||
(id, instance_name, provider_id, api_key, status, create_time, create_date, update_time, update_date)
|
||||
VALUES {', '.join(values)}
|
||||
"""
|
||||
self.db.execute_sql(insert_sql)
|
||||
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_instance table"""
|
||||
create_sql = """
|
||||
CREATE TABLE IF NOT EXISTS tenant_model_instance (
|
||||
id VARCHAR(32) NOT NULL PRIMARY KEY,
|
||||
instance_name VARCHAR(128) NOT NULL,
|
||||
provider_id VARCHAR(32) NOT NULL,
|
||||
api_key VARCHAR(512) NOT NULL,
|
||||
status VARCHAR(32) DEFAULT 'active',
|
||||
create_time BIGINT,
|
||||
create_date DATETIME,
|
||||
update_time BIGINT,
|
||||
update_date DATETIME,
|
||||
UNIQUE INDEX idx_api_key_provider_id (api_key, provider_id),
|
||||
INDEX idx_provider_id (provider_id)
|
||||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
||||
"""
|
||||
self.db.execute_sql(create_sql)
|
||||
logger.info("Created tenant_model_instance table")
|
||||
|
||||
|
||||
class TenantModelStage(MigrationStage):
|
||||
"""Migrate tenant_llm to tenant_model"""
|
||||
|
||||
name = "tenant_model"
|
||||
description = "Migrate tenant_llm to tenant_model (only status='0' records)"
|
||||
source_tables = ["tenant_llm", "tenant_model_provider", "tenant_model_instance"]
|
||||
target_tables = ["tenant_model"]
|
||||
|
||||
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 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.")
|
||||
return False
|
||||
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.")
|
||||
return False
|
||||
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.")
|
||||
return False
|
||||
logger.info("Target table 'tenant_model' does not exist, will create")
|
||||
return True
|
||||
|
||||
# Check if there's data to migrate (only status='0' records)
|
||||
cursor = self.db.execute_sql(
|
||||
"SELECT COUNT(*) FROM ("
|
||||
" SELECT tl.id "
|
||||
" FROM tenant_llm tl "
|
||||
" INNER JOIN tenant_model_provider tmp ON tmp.tenant_id = tl.tenant_id AND tmp.provider_name = tl.llm_factory "
|
||||
" INNER JOIN tenant_model_instance tmi ON tmi.provider_id = tmp.id AND tmi.api_key = tl.api_key "
|
||||
" WHERE tl.status = '0' "
|
||||
" AND NOT EXISTS ("
|
||||
" SELECT 1 FROM tenant_model tm "
|
||||
" WHERE tm.provider_id = tmp.id AND tm.model_name = tl.llm_name AND tm.instance_id = tmi.id"
|
||||
" )"
|
||||
") AS distinct_records"
|
||||
)
|
||||
count = cursor.fetchone()[0]
|
||||
|
||||
if count == 0:
|
||||
logger.info("No new data to migrate from tenant_llm to tenant_model (status='0' only)")
|
||||
return False
|
||||
|
||||
logger.info(f"Found {count} rows to migrate from tenant_llm to tenant_model")
|
||||
return True
|
||||
|
||||
def execute(self) -> tuple[int, list]:
|
||||
"""Execute migration"""
|
||||
current_ts = self.current_timestamp()
|
||||
rows_inserted = 0
|
||||
|
||||
# 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.")
|
||||
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.")
|
||||
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.")
|
||||
return 0, []
|
||||
logger.info("Target table 'tenant_model' 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 records from tenant_llm with provider_id and instance_id lookup
|
||||
# Only migrate records where status='0'
|
||||
cursor = self.db.execute_sql(
|
||||
"SELECT tl.id, tl.llm_name, tmp.id as provider_id, tmi.id as instance_id, "
|
||||
" tl.model_type, tl.status "
|
||||
"FROM tenant_llm tl "
|
||||
"INNER JOIN tenant_model_provider tmp ON tmp.tenant_id = tl.tenant_id AND tmp.provider_name = tl.llm_factory "
|
||||
"INNER JOIN tenant_model_instance tmi ON tmi.provider_id = tmp.id AND tmi.api_key = tl.api_key "
|
||||
"WHERE tl.status = '0' "
|
||||
"AND NOT EXISTS ("
|
||||
" SELECT 1 FROM tenant_model tm "
|
||||
" WHERE tm.provider_id = tmp.id AND tm.model_name = tl.llm_name AND tm.instance_id = tmi.id"
|
||||
")"
|
||||
)
|
||||
|
||||
records = cursor.fetchall()
|
||||
|
||||
if not records:
|
||||
logger.info("No records to migrate")
|
||||
return 0, []
|
||||
|
||||
logger.info(f"Migrating {len(records)} tenant_model records...")
|
||||
|
||||
if self.dry_run:
|
||||
logger.info(f"[DRY RUN] Would insert {len(records)} records")
|
||||
for source_id, llm_name, provider_id, instance_id, model_type, status in records[:5]:
|
||||
logger.info(f" model_name={llm_name}, provider_id={provider_id}, "
|
||||
f"instance_id={instance_id}, model_type={model_type}")
|
||||
if len(records) > 5:
|
||||
logger.info(f" ... and {len(records) - 5} more records")
|
||||
return len(records), self.target_tables
|
||||
|
||||
# Insert records in batches
|
||||
batch_size = 100
|
||||
for i in range(0, len(records), batch_size):
|
||||
batch = records[i:i + batch_size]
|
||||
values = []
|
||||
for source_id, llm_name, provider_id, instance_id, model_type, status in batch:
|
||||
record_id = self.generate_uuid()
|
||||
model_name_escaped = llm_name.replace("'", "''") if llm_name else ""
|
||||
model_type_escaped = model_type.replace("'", "''") if model_type else ""
|
||||
status_val = status if status else "active"
|
||||
values.append(f"('{record_id}', '{model_name_escaped}', '{provider_id}', "
|
||||
f"'{instance_id}', '{model_type_escaped}', '{status_val}', "
|
||||
f"{current_ts}, FROM_UNIXTIME({current_ts}), "
|
||||
f"{current_ts}, FROM_UNIXTIME({current_ts}))")
|
||||
|
||||
insert_sql = f"""
|
||||
INSERT INTO tenant_model
|
||||
(id, model_name, provider_id, instance_id, model_type, status,
|
||||
create_time, create_date, update_time, update_date)
|
||||
VALUES {', '.join(values)}
|
||||
"""
|
||||
self.db.execute_sql(insert_sql)
|
||||
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 table"""
|
||||
create_sql = """
|
||||
CREATE TABLE IF NOT EXISTS tenant_model (
|
||||
id VARCHAR(32) NOT NULL PRIMARY KEY,
|
||||
model_name VARCHAR(128),
|
||||
provider_id VARCHAR(32) NOT NULL,
|
||||
instance_id VARCHAR(32) NOT NULL,
|
||||
model_type VARCHAR(32) NOT NULL,
|
||||
status VARCHAR(32) DEFAULT 'active',
|
||||
create_time BIGINT,
|
||||
create_date DATETIME,
|
||||
update_time BIGINT,
|
||||
update_date DATETIME,
|
||||
INDEX idx_instance_id (instance_id),
|
||||
UNIQUE INDEX idx_provider_model_instance (provider_id, model_name, instance_id)
|
||||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
||||
"""
|
||||
self.db.execute_sql(create_sql)
|
||||
logger.info("Created tenant_model table")
|
||||
|
||||
|
||||
# Registry of available migration stages
|
||||
MIGRATION_STAGES = {
|
||||
'tenant_model_provider': TenantModelProviderStage,
|
||||
'tenant_model_instance': TenantModelInstanceStage,
|
||||
'tenant_model': TenantModelStage,
|
||||
}
|
||||
|
||||
|
||||
@@ -394,7 +747,7 @@ def run_migration(config: MigrationConfig, stages: list, dry_run: bool = True,
|
||||
total_stages = len(stages)
|
||||
|
||||
for idx, stage_name in enumerate(stages, 1):
|
||||
logger.info(f"\n{'=' * 60}")
|
||||
logger.info(f"{'=' * 60}")
|
||||
logger.info(f"Stage [{idx}/{total_stages}]: {stage_name}")
|
||||
logger.info(f"{'=' * 60}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user