mappings+migrate
This commit is contained in:
@@ -87,34 +87,72 @@ class MigrationPlugin(PluginBase):
|
||||
}
|
||||
|
||||
async def execute(self, params: Dict[str, Any]):
|
||||
from_env = params["from_env"]
|
||||
to_env = params["to_env"]
|
||||
dashboard_regex = params["dashboard_regex"]
|
||||
source_env_id = params.get("source_env_id")
|
||||
target_env_id = params.get("target_env_id")
|
||||
selected_ids = params.get("selected_ids")
|
||||
|
||||
# Legacy support or alternative params
|
||||
from_env_name = params.get("from_env")
|
||||
to_env_name = params.get("to_env")
|
||||
dashboard_regex = params.get("dashboard_regex")
|
||||
|
||||
replace_db_config = params.get("replace_db_config", False)
|
||||
from_db_id = params.get("from_db_id")
|
||||
to_db_id = params.get("to_db_id")
|
||||
|
||||
logger = SupersetLogger(log_dir=Path.cwd() / "logs", console=True)
|
||||
logger.info(f"[MigrationPlugin][Entry] Starting migration from {from_env} to {to_env}.")
|
||||
logger.info(f"[MigrationPlugin][Entry] Starting migration task.")
|
||||
logger.info(f"[MigrationPlugin][Action] Params: {params}")
|
||||
|
||||
try:
|
||||
config_manager = get_config_manager()
|
||||
all_clients = setup_clients(logger, custom_envs=config_manager.get_environments())
|
||||
from_c = all_clients.get(from_env)
|
||||
to_c = all_clients.get(to_env)
|
||||
environments = config_manager.get_environments()
|
||||
|
||||
# Resolve environments
|
||||
src_env = None
|
||||
tgt_env = None
|
||||
|
||||
if source_env_id:
|
||||
src_env = next((e for e in environments if e.id == source_env_id), None)
|
||||
elif from_env_name:
|
||||
src_env = next((e for e in environments if e.name == from_env_name), None)
|
||||
|
||||
if target_env_id:
|
||||
tgt_env = next((e for e in environments if e.id == target_env_id), None)
|
||||
elif to_env_name:
|
||||
tgt_env = next((e for e in environments if e.name == to_env_name), None)
|
||||
|
||||
if not src_env or not tgt_env:
|
||||
raise ValueError(f"Could not resolve source or target environment. Source: {source_env_id or from_env_name}, Target: {target_env_id or to_env_name}")
|
||||
|
||||
from_env_name = src_env.name
|
||||
to_env_name = tgt_env.name
|
||||
|
||||
logger.info(f"[MigrationPlugin][State] Resolved environments: {from_env_name} -> {to_env_name}")
|
||||
|
||||
all_clients = setup_clients(logger, custom_envs=environments)
|
||||
from_c = all_clients.get(from_env_name)
|
||||
to_c = all_clients.get(to_env_name)
|
||||
|
||||
if not from_c or not to_c:
|
||||
raise ValueError(f"One or both environments ('{from_env}', '{to_env}') not found in configuration.")
|
||||
raise ValueError(f"Clients not initialized for environments: {from_env_name}, {to_env_name}")
|
||||
|
||||
_, all_dashboards = from_c.get_dashboards()
|
||||
|
||||
regex_str = str(dashboard_regex)
|
||||
dashboards_to_migrate = [
|
||||
d for d in all_dashboards if re.search(regex_str, d["dashboard_title"], re.IGNORECASE)
|
||||
]
|
||||
dashboards_to_migrate = []
|
||||
if selected_ids:
|
||||
dashboards_to_migrate = [d for d in all_dashboards if d["id"] in selected_ids]
|
||||
elif dashboard_regex:
|
||||
regex_str = str(dashboard_regex)
|
||||
dashboards_to_migrate = [
|
||||
d for d in all_dashboards if re.search(regex_str, d["dashboard_title"], re.IGNORECASE)
|
||||
]
|
||||
else:
|
||||
logger.warning("[MigrationPlugin][State] No selection criteria provided (selected_ids or dashboard_regex).")
|
||||
return
|
||||
|
||||
if not dashboards_to_migrate:
|
||||
logger.warning("[MigrationPlugin][State] No dashboards found matching the regex.")
|
||||
logger.warning("[MigrationPlugin][State] No dashboards found matching criteria.")
|
||||
return
|
||||
|
||||
# Fetch mappings from database
|
||||
@@ -123,8 +161,8 @@ class MigrationPlugin(PluginBase):
|
||||
db = SessionLocal()
|
||||
try:
|
||||
# Find environment IDs by name
|
||||
src_env = db.query(Environment).filter(Environment.name == from_env).first()
|
||||
tgt_env = db.query(Environment).filter(Environment.name == to_env).first()
|
||||
src_env = db.query(Environment).filter(Environment.name == from_env_name).first()
|
||||
tgt_env = db.query(Environment).filter(Environment.name == to_env_name).first()
|
||||
|
||||
if src_env and tgt_env:
|
||||
mappings = db.query(DatabaseMapping).filter(
|
||||
@@ -144,46 +182,39 @@ class MigrationPlugin(PluginBase):
|
||||
try:
|
||||
exported_content, _ = from_c.export_dashboard(dash_id)
|
||||
with create_temp_file(content=exported_content, dry_run=True, suffix=".zip", logger=logger) as tmp_zip_path:
|
||||
if not replace_db_config:
|
||||
to_c.import_dashboard(file_name=tmp_zip_path, dash_id=dash_id, dash_slug=dash_slug)
|
||||
else:
|
||||
# Check for missing mappings before transformation
|
||||
# This is a simplified check, in reality we'd check all YAMLs
|
||||
# For US3, we'll just use the engine and handle missing ones there
|
||||
with create_temp_file(suffix=".zip", dry_run=True, logger=logger) as tmp_new_zip:
|
||||
# If we have missing mappings, we might need to pause
|
||||
# For now, let's assume the engine can tell us what's missing
|
||||
success = engine.transform_zip(str(tmp_zip_path), str(tmp_new_zip), db_mapping)
|
||||
|
||||
if not success:
|
||||
# Signal missing mapping and wait
|
||||
task_id = params.get("_task_id")
|
||||
if task_id:
|
||||
from ..dependencies import get_task_manager
|
||||
tm = get_task_manager()
|
||||
logger.info(f"[MigrationPlugin][Action] Pausing for missing mapping in task {task_id}")
|
||||
# In a real scenario, we'd pass the missing DB info to the frontend
|
||||
# For this task, we'll just simulate the wait
|
||||
await tm.wait_for_resolution(task_id)
|
||||
# After resolution, retry transformation with updated mappings
|
||||
# (Mappings would be updated in task.params by resolve_task)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
src_env = db.query(Environment).filter(Environment.name == from_env).first()
|
||||
tgt_env = db.query(Environment).filter(Environment.name == to_env).first()
|
||||
mappings = db.query(DatabaseMapping).filter(
|
||||
DatabaseMapping.source_env_id == src_env.id,
|
||||
DatabaseMapping.target_env_id == tgt_env.id
|
||||
).all()
|
||||
db_mapping = {m.source_db_uuid: m.target_db_uuid for m in mappings}
|
||||
finally:
|
||||
db.close()
|
||||
success = engine.transform_zip(str(tmp_zip_path), str(tmp_new_zip), db_mapping)
|
||||
# Always transform to strip databases to avoid password errors
|
||||
with create_temp_file(suffix=".zip", dry_run=True, logger=logger) as tmp_new_zip:
|
||||
success = engine.transform_zip(str(tmp_zip_path), str(tmp_new_zip), db_mapping, strip_databases=False)
|
||||
|
||||
if not success and replace_db_config:
|
||||
# Signal missing mapping and wait (only if we care about mappings)
|
||||
task_id = params.get("_task_id")
|
||||
if task_id:
|
||||
from ..dependencies import get_task_manager
|
||||
tm = get_task_manager()
|
||||
logger.info(f"[MigrationPlugin][Action] Pausing for missing mapping in task {task_id}")
|
||||
# In a real scenario, we'd pass the missing DB info to the frontend
|
||||
# For this task, we'll just simulate the wait
|
||||
await tm.wait_for_resolution(task_id)
|
||||
# After resolution, retry transformation with updated mappings
|
||||
# (Mappings would be updated in task.params by resolve_task)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
src_env = db.query(Environment).filter(Environment.name == from_env_name).first()
|
||||
tgt_env = db.query(Environment).filter(Environment.name == to_env_name).first()
|
||||
mappings = db.query(DatabaseMapping).filter(
|
||||
DatabaseMapping.source_env_id == src_env.id,
|
||||
DatabaseMapping.target_env_id == tgt_env.id
|
||||
).all()
|
||||
db_mapping = {m.source_db_uuid: m.target_db_uuid for m in mappings}
|
||||
finally:
|
||||
db.close()
|
||||
success = engine.transform_zip(str(tmp_zip_path), str(tmp_new_zip), db_mapping, strip_databases=False)
|
||||
|
||||
if success:
|
||||
to_c.import_dashboard(file_name=tmp_new_zip, dash_id=dash_id, dash_slug=dash_slug)
|
||||
else:
|
||||
logger.error(f"[MigrationPlugin][Failure] Failed to transform ZIP for dashboard {title}")
|
||||
if success:
|
||||
to_c.import_dashboard(file_name=tmp_new_zip, dash_id=dash_id, dash_slug=dash_slug)
|
||||
else:
|
||||
logger.error(f"[MigrationPlugin][Failure] Failed to transform ZIP for dashboard {title}")
|
||||
|
||||
logger.info(f"[MigrationPlugin][Success] Dashboard {title} imported.")
|
||||
except Exception as exc:
|
||||
|
||||
Reference in New Issue
Block a user