Spaces:

arcticaurora
/

Migrate

Sleeping

App Files Files Community

arcticaurora commited on Sep 27, 2025

Commit

6639025

verified ·

1 Parent(s): f4c0131

Update app.py

Browse files

Files changed (1) hide show

app.py +234 -411

app.py CHANGED Viewed

@@ -12,12 +12,12 @@ import time
 import json
 import datetime
 import uuid
-import shutil
-from typing import Dict, Any, Optional, List
 from pathlib import Path
 import psycopg2
 import logging
 import signal
 # --- Setup Logging ---
 logging.basicConfig(
@@ -44,7 +44,6 @@ os.makedirs("static/css", exist_ok=True)
 os.makedirs("static/js", exist_ok=True)
 os.makedirs("dumps", exist_ok=True)
-# Create placeholder files if they don't exist so the app can run
 if not os.path.exists("templates/index.html"):
     with open("templates/index.html", "w") as f:
         f.write("<!DOCTYPE html><html><head><title>Migrator</title></head><body>Loading...</body></html>")
@@ -54,12 +53,11 @@ templates = Jinja2Templates(directory="templates")
 app.mount("/static", StaticFiles(directory="static"), name="static")
 app.mount("/downloads", StaticFiles(directory="dumps"), name="downloads")
 # --- Global State for Migration ---
 migration_state = {
     "id": str(uuid.uuid4()),
     "running": False,
-    "operation": None,  # "dump" or "restore"
     "start_time": None,
     "end_time": None,
     "dump_file": None,
@@ -77,277 +75,168 @@ migration_state = {
         "size_processed_bytes": 0,
         "percent_complete_by_size": 0,
         "percent_complete_by_count": 0,
-        "manifest": None, # Will store the detailed chunk manifest
-        # Legacy fields for non-TimescaleDB or fallback
-        "tables_completed": 0,
-        "total_tables": 0,
-        "current_size_mb": 0,
-        "growth_rate_mb_per_sec": 0,
     }
 }
-# Lock for updating global state
 migration_lock = threading.Lock()
-def log_message(message: str, level: str = "info", command: str = None):
-    """Add timestamped log message with level"""
     timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    log_entry = {
-        "timestamp": timestamp,
-        "message": message,
-        "level": level,
-        "command": command,
-        "id": len(migration_state["log"])
-    }
     with migration_lock:
         migration_state["log"].append(log_entry)
         migration_state["last_activity"] = time.time()
     logger.info(f"[{level.upper()}] {message}")
-    if command:
-        logger.info(f"Command: {command}")
 def test_connection_logic(connection_string: str) -> bool:
-    """Test a PostgreSQL connection string (internal logic)"""
     try:
-        conn = psycopg2.connect(connection_string)
-        conn.close()
-        return True
     except Exception as e:
-        logger.error(f"Connection test failed: {str(e)}")
         return False
 def get_chunk_manifest(connection_string: str) -> Dict[str, Any]:
-    """
-    Connects to a TimescaleDB and generates a manifest of all hypertable chunks
-    (both regular and compressed) and their sizes.
-    """
-    manifest = {
-        "total_chunks": 0,
-        "total_size_bytes": 0,
-        "chunks": {} # Format: {'schema.chunk_name': size_in_bytes}
-    }
-    log_message("Attempting to generate chunk manifest from source database...", "info")
     try:
-        conn = psycopg2.connect(connection_string)
-        with conn.cursor() as cur:
-            # This query joins chunks with their respective hypertables to get a full list
-            # of all chunk tables (including compressed ones) that pg_dump will process.
-            query = """
-                SELECT
-                    c.schema_name || '.' || c.table_name as chunk_fqn,
-                    pg_total_relation_size(c.schema_name || '.' || c.table_name) as chunk_size_bytes
-                FROM _timescaledb_catalog.chunk c;
-            """
-            cur.execute(query)
-            results = cur.fetchall()
-            if not results:
-                log_message("No TimescaleDB chunks found. Proceeding with standard progress tracking.", "warning")
-                return manifest
-            for row in results:
-                chunk_name, chunk_size = row
-                if chunk_size is None: continue # Skip if size is null for any reason
-                manifest["chunks"][chunk_name] = chunk_size
-                manifest["total_size_bytes"] += chunk_size
-            manifest["total_chunks"] = len(manifest["chunks"])
-            log_message(
-                f"Successfully generated manifest: {manifest['total_chunks']} chunks found, "
-                f"total size {round(manifest['total_size_bytes'] / (1024**3), 2)} GB.",
-                "success"
-            )
-            return manifest
     except Exception as e:
         log_message(f"Failed to generate chunk manifest: {e}. Will use basic progress tracking.", "error")
-        # Return an empty manifest on failure so the process can continue
-        return { "total_chunks": 0, "total_size_bytes": 0, "chunks": {} }
-    finally:
-        if 'conn' in locals() and conn:
-            conn.close()
-def run_dump(source_conn: str, file_path: str, options: dict):
-    """Run pg_dump in a background thread with enhanced progress tracking."""
-    try:
-        absolute_file_path = os.path.abspath(file_path)
-        if os.path.exists(absolute_file_path):
-            os.remove(absolute_file_path)
-        env = os.environ.copy()
-        format_flag = "-F" + options.get("format", "c")
-        cmd = ["pg_dump", source_conn, format_flag, "-v", "-f", absolute_file_path]
-        if options.get("schema"):
-            cmd.extend(["-n", options["schema"]])
-        if options.get("compression") and options["compression"] != "default":
-            cmd.extend(["-Z", options["compression"]])
-        log_message(f"Starting database dump to {absolute_file_path}", "info", " ".join(cmd))
-        with migration_lock:
-            migration_state["start_time"] = time.time()
-            migration_state["running"] = True
-            migration_state["operation"] = "dump"
-            migration_state["dump_file"] = absolute_file_path
-            migration_state["dump_completed"] = False
-        process = subprocess.Popen(
-            cmd,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            env=env,
-            text=True,
-            bufsize=1
-        )
         with migration_lock:
             migration_state["process"] = process
-        def read_stderr():
-            """Reads stderr from pg_dump and updates progress based on the manifest."""
-            for line in iter(process.stderr.readline, ''):
-                line = line.strip()
-                if not line:
-                    continue
-                log_message(line, "info") # Log all verbose output
-                # Check for table data dumping lines
-                if "dumping data for table" in line:
-                    try:
-                        # Extract the fully-qualified table/chunk name
-                        parts = line.split(" for table ")[1].split(" ")
-                        table_name_full = parts[0]
-                        with migration_lock:
-                            migration_state["progress"]["current_table"] = table_name_full
-                            manifest = migration_state["progress"].get("manifest")
-                            # Use manifest for accurate progress if available
-                            if manifest and manifest.get("chunks"):
-                                chunk_size = manifest["chunks"].get(table_name_full)
-                                if chunk_size is not None: # It's a chunk we're tracking
-                                    migration_state["progress"]["chunks_processed"] += 1
-                                    migration_state["progress"]["size_processed_bytes"] += chunk_size
-                                    # Recalculate percentages
-                                    if migration_state["progress"]["total_chunks"] > 0:
-                                        migration_state["progress"]["percent_complete_by_count"] = round(
-                                            (migration_state["progress"]["chunks_processed"] / migration_state["progress"]["total_chunks"]) * 100, 2
-                                        )
-                                    if migration_state["progress"]["total_size_bytes"] > 0:
-                                        migration_state["progress"]["percent_complete_by_size"] = round(
-                                            (migration_state["progress"]["size_processed_bytes"] / migration_state["progress"]["total_size_bytes"]) * 100, 2
-                                        )
-                                else:
-                                    # It's a regular table, not a chunk
-                                    migration_state["progress"]["tables_completed"] += 1
-                            else:
-                                # Fallback to simple table counting
-                                migration_state["progress"]["tables_completed"] += 1
-                    except Exception as e:
-                        log_message(f"Error parsing pg_dump output line '{line}': {e}", "warning")
-        stderr_thread = threading.Thread(target=read_stderr, daemon=True)
         stderr_thread.start()
         exit_code = process.wait()
         stderr_thread.join(timeout=2.0)
-        final_size = 0
-        if os.path.exists(absolute_file_path):
-            final_size = os.path.getsize(absolute_file_path)
         with migration_lock:
             migration_state["dump_file_size"] = final_size / (1024 * 1024) # MB
             migration_state["end_time"] = time.time()
             migration_state["running"] = False
             migration_state["process"] = None
-            if exit_code == 0 and final_size > 0:
-                migration_state["dump_completed"] = True
                 total_time = migration_state["end_time"] - migration_state["start_time"]
-                log_message(
-                    f"Database dump completed successfully. Size: {migration_state['dump_file_size']:.2f} MB. Time: {total_time:.2f} seconds",
-                    "success"
-                )
             else:
-                migration_state["dump_completed"] = False
-                if exit_code != 0:
-                    log_message(f"Database dump failed with exit code {exit_code}", "error")
-                elif final_size == 0:
-                    log_message(f"Dump file was created but is empty: {absolute_file_path}", "error")
     except Exception as e:
-        log_message(f"An unexpected error occurred during database dump: {str(e)}", "error")
         with migration_lock:
             migration_state["running"] = False
             migration_state["process"] = None
 def run_restore(target_conn: str, file_path: str, options: dict):
-    """Run pg_restore in a background thread."""
     try:
-        if not os.path.exists(file_path):
-            log_message(f"Dump file not found: {file_path}", "error")
-            return
-        env = os.environ.copy()
         # Pre-restore step
         if options.get("timescaledb_pre_restore", True):
             pre_restore_cmd = ["psql", target_conn, "-c", "SELECT timescaledb_pre_restore();"]
-            log_message("Running timescaledb_pre_restore()", "info", " ".join(pre_restore_cmd))
-            pre_restore_proc = subprocess.run(pre_restore_cmd, capture_output=True, text=True, env=env)
             if pre_restore_proc.returncode != 0:
                 log_message(f"Pre-restore failed: {pre_restore_proc.stderr or pre_restore_proc.stdout}", "error")
                 return
-        # Build pg_restore command
         cmd = ["pg_restore", "-d", target_conn, "-v"]
         if options.get("no_owner", True): cmd.append("--no-owner")
         if options.get("clean", False): cmd.append("--clean")
-        # The single-transaction option is removed as per user request.
-        # if options.get("single_transaction", True): cmd.append("--single-transaction")
         cmd.append(file_path)
-        log_message(f"Starting database restore from {file_path}", "info", " ".join(cmd))
-        with migration_lock:
-            migration_state["start_time"] = time.time()
-            migration_state["running"] = True
-            migration_state["operation"] = "restore"
-            migration_state["restore_completed"] = False
-        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env, text=True, bufsize=1)
         with migration_lock:
             migration_state["process"] = process
-        # For restore, we use the simpler table counting as we don't have a pre-restore manifest
-        def read_stderr():
-            for line in iter(process.stderr.readline, ''):
-                line = line.strip()
-                if not line: continue
-                log_message(line, "info")
-                if "processing" in line.lower() and "table data" in line.lower():
-                    try:
-                        table_name = line.split("table data ")[1].split(" ")[0]
-                        with migration_lock:
-                            migration_state["progress"]["current_table"] = table_name
-                            migration_state["progress"]["tables_completed"] += 1
-                    except Exception as e:
-                        logger.warning(f"Could not parse table name from restore line: {line}, Error: {e}")
-        stderr_thread = threading.Thread(target=read_stderr, daemon=True)
         stderr_thread.start()
         exit_code = process.wait()
         stderr_thread.join(timeout=2.0)
         post_restore_success = True
         if exit_code == 0 and options.get("timescaledb_post_restore", True):
             post_restore_cmd = ["psql", target_conn, "-c", "SELECT timescaledb_post_restore(); ANALYZE;"]
-            log_message("Running timescaledb_post_restore() and ANALYZE", "info", " ".join(post_restore_cmd))
-            post_restore_proc = subprocess.run(post_restore_cmd, capture_output=True, text=True, env=env)
             if post_restore_proc.returncode != 0:
                 log_message(f"Post-restore failed: {post_restore_proc.stderr or post_restore_proc.stdout}", "error")
                 post_restore_success = False
@@ -356,285 +245,219 @@ def run_restore(target_conn: str, file_path: str, options: dict):
             migration_state["end_time"] = time.time()
             migration_state["running"] = False
             migration_state["process"] = None
-            if exit_code == 0 and post_restore_success:
-                migration_state["restore_completed"] = True
                 total_time = migration_state["end_time"] - migration_state["start_time"]
-                log_message(f"Database restore completed successfully. Time: {total_time:.2f} seconds", "success")
             else:
-                migration_state["restore_completed"] = False
-                if exit_code != 0:
-                    log_message(f"Database restore failed with exit code {exit_code}", "error")
     except Exception as e:
-        log_message(f"An unexpected error occurred during database restore: {e}", "error")
         with migration_lock:
             migration_state["running"] = False
             migration_state["process"] = None
 def stop_current_process():
-    """Stop the current running process forcefully."""
     with migration_lock:
         process = migration_state.get("process")
         if not process or not migration_state.get("running"):
-            log_message("Stop command received, but no process was running.", "info")
-            if migration_state["running"]: # Correct inconsistent state
-                migration_state["running"] = False
             return False
         try:
             pid = process.pid
-            operation = migration_state["operation"]
-            log_message(f"Attempting to stop {operation} process (PID: {pid})...", "warning")
-            if process.poll() is not None:
-                log_message(f"Process (PID: {pid}) already terminated.", "info")
             else:
-                # Use os.killpg to terminate the entire process group, which is more reliable
-                # for stopping child processes spawned by shell commands.
-                if hasattr(os, 'killpg') and hasattr(os, 'getpgid'):
-                    try:
-                        os.killpg(os.getpgid(pid), signal.SIGKILL)
-                        log_message(f"Sent SIGKILL to process group of PID {pid}.", "warning")
-                    except ProcessLookupError:
-                        log_message(f"Process group for PID {pid} not found, it likely terminated.", "info")
-                    except Exception as e:
-                        log_message(f"Failed to kill process group for PID {pid}: {e}. Falling back to process.kill().", "error")
-                        process.kill() # Fallback for other errors or non-Unix systems
-                else:
-                    process.kill()
-                    log_message(f"Sent SIGKILL directly to process (PID: {pid}).", "warning")
-                # Wait a moment for the OS to process the kill signal
-                process.wait(timeout=5)
-            log_message(f"Database {operation} operation stopped by user.", "warning")
             return True
-        except (ProcessLookupError, subprocess.TimeoutExpired):
-             log_message(f"Process was already gone or did not terminate in time.", "info")
-             return True # Consider it a success as the process is no longer running
         except Exception as e:
             log_message(f"Error stopping process: {e}", "error")
             return False
         finally:
-            # Always update state to reflect the stop action
             migration_state["running"] = False
             migration_state["process"] = None
             migration_state["end_time"] = time.time()
 # --- API Endpoints ---
 @app.get("/", response_class=HTMLResponse)
 async def home(request: Request):
-    """Serve the main HTML file."""
     return templates.TemplateResponse("index.html", {"request": request})
 @app.post("/start-dump")
 async def start_dump_endpoint(data: Dict[str, Any], background_tasks: BackgroundTasks):
-    """Start a database dump process after generating a chunk manifest."""
     source_conn = data.get("source_conn")
     options = data.get("options", {})
-    if not source_conn:
-        raise HTTPException(status_code=400, detail="Source connection string is required")
-    if not test_connection_logic(source_conn):
-        raise HTTPException(status_code=400, detail="Source connection failed. Cannot start dump.")
     if migration_state["running"]:
-        log_message("Stop command issued to clear running process for new dump.", "warning")
         stop_current_process()
-        time.sleep(1) # Give a moment for the process to die
-    # Generate the detailed chunk manifest before starting
-    manifest = get_chunk_manifest(source_conn)
-    # Setup file path
-    filename = options.get("filename", "timescale_backup").strip().replace(" ", "_")
-    filename = "".join(c for c in filename if c.isalnum() or c in ('_','-'))
-    if not filename: filename = "timescale_backup"
-    format_flag = options.get("format", "c")
-    extension = ".dump"
-    if format_flag == "p": extension = ".sql"
-    elif format_flag == "d": extension = ""
-    elif format_flag == "t": extension = ".tar"
-    dumps_dir = Path("dumps").resolve()
-    file_path = dumps_dir / f"{filename}{extension}"
-    if not str(file_path).startswith(str(dumps_dir)):
-        raise HTTPException(status_code=400, detail="Invalid filename specified.")
-    # Reset state with manifest data
     with migration_lock:
-        migration_state["id"] = str(uuid.uuid4())
-        migration_state["operation"] = "dump"
-        migration_state["dump_file"] = str(file_path)
-        migration_state["log"] = [] # Clear logs for new operation
-        migration_state["progress"] = {
-            "current_table": None,
-            "total_chunks": manifest["total_chunks"],
-            "chunks_processed": 0,
-            "total_size_bytes": manifest["total_size_bytes"],
-            "size_processed_bytes": 0,
-            "percent_complete_by_size": 0,
-            "percent_complete_by_count": 0,
-            "manifest": manifest,
-            "tables_completed": 0,
-            "total_tables": 0,
-            "current_size_mb": 0,
-            "growth_rate_mb_per_sec": 0,
-        }
-        # Reset other state variables
-        migration_state["running"] = False
-        migration_state["start_time"] = None
-        migration_state["end_time"] = None
-        migration_state["dump_completed"] = False
-        migration_state["restore_completed"] = False
-    background_tasks.add_task(run_dump, source_conn, str(file_path), options)
-    return JSONResponse(content={"success": True, "message": "Dump process initiated"})
 @app.post("/start-restore")
 async def start_restore_endpoint(data: Dict[str, Any], background_tasks: BackgroundTasks):
-    """Start a database restore process."""
     target_conn = data.get("target_conn")
     dump_file = data.get("dump_file")
     options = data.get("options", {})
-    if not target_conn or not dump_file:
-        raise HTTPException(status_code=400, detail="Target connection and dump file are required.")
-    if not test_connection_logic(target_conn):
-        raise HTTPException(status_code=400, detail="Target connection failed. Cannot start restore.")
-    dumps_dir = Path("dumps").resolve()
-    dump_file_path = Path(dump_file).resolve()
-    if not dump_file_path.exists() or not str(dump_file_path).startswith(str(dumps_dir)):
-        raise HTTPException(status_code=400, detail="Invalid or non-existent dump file selected.")
     if migration_state["running"]:
-        log_message("Stop command issued to clear running process for new restore.", "warning")
         stop_current_process()
         time.sleep(1)
-    # Reset state for restore operation
     with migration_lock:
-        migration_state["id"] = str(uuid.uuid4())
-        migration_state["operation"] = "restore"
-        migration_state["log"] = [] # Clear logs
-        migration_state["progress"] = { # Reset to defaults for restore
-            "current_table": None, "total_chunks": 0, "chunks_processed": 0,
-            "total_size_bytes": 0, "size_processed_bytes": 0, "percent_complete_by_size": 0,
-            "percent_complete_by_count": 0, "manifest": None, "tables_completed": 0,
-            "total_tables": 0, "current_size_mb": 0, "growth_rate_mb_per_sec": 0,
-        }
-        migration_state["running"] = False
-        migration_state["start_time"] = None
-        migration_state["end_time"] = None
-        migration_state["dump_completed"] = False
-        migration_state["restore_completed"] = False
-    background_tasks.add_task(run_restore, target_conn, str(dump_file_path), options)
-    return JSONResponse(content={"success": True, "message": "Restore process initiated"})
 @app.post("/stop-process")
 async def stop_process_endpoint():
-    """Stop the current database process."""
-    stopped = stop_current_process()
-    if stopped:
-        return JSONResponse(content={"success": True, "message": "Process stop initiated."})
-    else:
-        # Check if it wasn't running or if stop failed
-        with migration_lock:
-            is_running = migration_state["running"]
-        if not is_running:
-            return JSONResponse(content={"success": False, "message": "No process was running to stop."})
-        else:
-            return JSONResponse(status_code=500, content={"success": False, "message": "Failed to stop the process. Check logs."})
 @app.get("/status")
 async def get_status():
-    """Get the current migration status."""
     with migration_lock:
-        # Create a deep copy to avoid race conditions and to remove non-serializable objects
         state_copy = {k: v for k, v in migration_state.items() if k != "process"}
-        # Don't send the full manifest back in every status update to save bandwidth
         if "progress" in state_copy and "manifest" in state_copy["progress"]:
-            state_copy["progress"]["manifest"] = None
     return JSONResponse(content=state_copy)
 @app.get("/list-dumps")
 async def list_dumps():
-    """List available dump files."""
-    try:
-        dumps_dir = Path("dumps")
-        dump_files = []
-        for f in dumps_dir.iterdir():
-            if f.is_file() or f.is_dir():
-                try:
-                    stat = f.stat()
-                    size = sum(file.stat().st_size for file in f.rglob('*')) if f.is_dir() else stat.st_size
-                    dump_files.append({
-                        "name": f.name,
-                        "path": str(f),
-                        "size_mb": size / (1024 * 1024),
-                        "date": datetime.datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %H:%M:%S"),
-                        "is_dir": f.is_dir()
-                    })
-                except OSError as e:
-                    logger.error(f"Could not stat file/dir {f.name}: {e}")
-        dump_files.sort(key=lambda x: x["date"], reverse=True)
-        return JSONResponse(content={"success": True, "dumps": dump_files})
-    except Exception as e:
-        logger.error(f"Failed to list dumps: {e}")
-        raise HTTPException(status_code=500, detail=f"Failed to list dumps: {e}")
 @app.get("/downloads/{file_name:path}")
 async def download_file(file_name: str):
-    """Download a dump file."""
-    dumps_dir = Path("dumps").resolve()
-    file_path = (dumps_dir / file_name).resolve()
-    if not file_path.exists() or not str(file_path).startswith(str(dumps_dir)):
-        raise HTTPException(status_code=404, detail="File not found or invalid path")
-    if file_path.is_dir():
-        raise HTTPException(status_code=400, detail="Directory downloads are not supported.")
-    return FileResponse(path=str(file_path), filename=file_name, media_type="application/octet-stream")
-# --- Other Endpoints (Test Connection, etc.) ---
-# These are largely unchanged but included for completeness.
 @app.post("/test-connection")
 async def test_connection_endpoint(data: Dict[str, str]):
-    """Test a database connection and get basic info"""
     connection_string = data.get("connection_string")
-    if not connection_string:
-        raise HTTPException(status_code=400, detail="Connection string is required")
-    if not test_connection_logic(connection_string):
-        return JSONResponse(content={"success": False, "message": "Failed to connect to database"})
     try:
-        conn = psycopg2.connect(connection_string)
-        with conn.cursor() as cur:
-            cur.execute("SELECT version()")
-            version = cur.fetchone()[0]
-            cur.execute("SELECT current_database()")
-            database = cur.fetchone()[0]
-            is_timescaledb, ts_version = False, None
-            cur.execute("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'timescaledb');")
-            if cur.fetchone()[0]:
-                is_timescaledb = True
-                cur.execute("SELECT extversion FROM pg_extension WHERE extname = 'timescaledb'")
-                ts_version = cur.fetchone()[0]
-        log_message(f"Successful connection test to database: {database}", "success")
-        return JSONResponse(content={
             "success": True, "version": version, "is_timescaledb": is_timescaledb,
             "timescaledb_version": ts_version, "database": database
         })
     except Exception as e:
-        log_message(f"Connection test failed during info fetch: {e}", "error")
         raise HTTPException(status_code=500, detail=f"Database error: {e}")
-    finally:
-        if 'conn' in locals() and conn:
-            conn.close()
 if __name__ == "__main__":
     import uvicorn

 import json
 import datetime
 import uuid
+from typing import Dict, Any
 from pathlib import Path
 import psycopg2
 import logging
 import signal
+import re
 # --- Setup Logging ---
 logging.basicConfig(
 os.makedirs("static/js", exist_ok=True)
 os.makedirs("dumps", exist_ok=True)
 if not os.path.exists("templates/index.html"):
     with open("templates/index.html", "w") as f:
         f.write("<!DOCTYPE html><html><head><title>Migrator</title></head><body>Loading...</body></html>")
 app.mount("/static", StaticFiles(directory="static"), name="static")
 app.mount("/downloads", StaticFiles(directory="dumps"), name="downloads")
 # --- Global State for Migration ---
 migration_state = {
     "id": str(uuid.uuid4()),
     "running": False,
+    "operation": None,
     "start_time": None,
     "end_time": None,
     "dump_file": None,
         "size_processed_bytes": 0,
         "percent_complete_by_size": 0,
         "percent_complete_by_count": 0,
+        "manifest_loaded": False,
+        "manifest": None, # Full manifest only stored during operation, not sent via status
+        "tables_completed": 0, # Fallback for non-manifest restores
     }
 }
 migration_lock = threading.Lock()
+# --- Helper Functions ---
+def log_message(message: str, level: str = "info"):
     timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    log_entry = {"timestamp": timestamp, "message": message, "level": level, "id": len(migration_state.get("log", []))}
     with migration_lock:
         migration_state["log"].append(log_entry)
         migration_state["last_activity"] = time.time()
     logger.info(f"[{level.upper()}] {message}")
 def test_connection_logic(connection_string: str) -> bool:
     try:
+        with psycopg2.connect(connection_string) as conn:
+            return True
     except Exception as e:
+        logger.error(f"Connection test failed: {e}")
         return False
 def get_chunk_manifest(connection_string: str) -> Dict[str, Any]:
+    manifest = {"total_chunks": 0, "total_size_bytes": 0, "chunks": {}}
+    log_message("Generating chunk manifest from source database...")
     try:
+        with psycopg2.connect(connection_string) as conn:
+            with conn.cursor() as cur:
+                query = """
+                    SELECT c.schema_name || '.' || c.table_name as chunk_fqn,
+                           pg_total_relation_size(c.schema_name || '.' || c.table_name) as chunk_size_bytes
+                    FROM _timescaledb_catalog.chunk c;
+                """
+                cur.execute(query)
+                for row in cur.fetchall():
+                    chunk_name, chunk_size = row
+                    if chunk_size is not None:
+                        manifest["chunks"][chunk_name] = chunk_size
+                        manifest["total_size_bytes"] += chunk_size
+                manifest["total_chunks"] = len(manifest["chunks"])
+        if manifest["total_chunks"] > 0:
+            log_message(f"Manifest generated: {manifest['total_chunks']} chunks, total size {round(manifest['total_size_bytes'] / (1024**3), 2)} GB.", "success")
+        else:
+            log_message("No TimescaleDB chunks found. Will use standard progress tracking.", "warning")
+        return manifest
     except Exception as e:
         log_message(f"Failed to generate chunk manifest: {e}. Will use basic progress tracking.", "error")
+        return {"total_chunks": 0, "total_size_bytes": 0, "chunks": {}}
+def process_stderr(process, operation):
+    """Generic stderr reader for dump and restore, supports manifest-based progress."""
+    for line in iter(process.stderr.readline, ''):
+        line = line.strip()
+        if not line: continue
+        log_message(line)
+        # Regex to handle various pg_dump/pg_restore verbose outputs for tables/chunks
+        match = re.search(r'(?:table|data for table|processing|from table)\s+"?([^".\s]+)"?\.?"?([^".\s]+)"?', line)
+        if match:
+            schema, table = match.groups()
+            table_name_full = f"{schema}.{table}"
+            with migration_lock:
+                migration_state["progress"]["current_table"] = table_name_full
+                manifest = migration_state["progress"].get("manifest")
+                if manifest and manifest.get("chunks"):
+                    chunk_size = manifest["chunks"].get(table_name_full)
+                    if chunk_size is not None:
+                        migration_state["progress"]["chunks_processed"] += 1
+                        migration_state["progress"]["size_processed_bytes"] += chunk_size
+                        if migration_state["progress"]["total_chunks"] > 0:
+                            migration_state["progress"]["percent_complete_by_count"] = round((migration_state["progress"]["chunks_processed"] / migration_state["progress"]["total_chunks"]) * 100, 2)
+                        if migration_state["progress"]["total_size_bytes"] > 0:
+                            migration_state["progress"]["percent_complete_by_size"] = round((migration_state["progress"]["size_processed_bytes"] / migration_state["progress"]["total_size_bytes"]) * 100, 2)
+                    else:
+                        migration_state["progress"]["tables_completed"] += 1 # A non-chunk table
+                else:
+                    migration_state["progress"]["tables_completed"] += 1 # Fallback mode
+# --- Core Background Tasks ---
+def run_dump(source_conn: str, file_path: str, manifest_path: str, manifest_data: dict, options: dict):
+    try:
+        cmd = ["pg_dump", source_conn, f"-F{options.get('format', 'c')}", "-v", "-f", file_path]
+        if options.get("schema"): cmd.extend(["-n", options["schema"]])
+        if options.get("compression", "default") != "default": cmd.extend(["-Z", options["compression"]])
+        log_message(f"Starting database dump...", "info")
+        process = subprocess.Popen(cmd, stderr=subprocess.PIPE, text=True, bufsize=1)
         with migration_lock:
             migration_state["process"] = process
+        stderr_thread = threading.Thread(target=process_stderr, args=(process, "dump"), daemon=True)
         stderr_thread.start()
         exit_code = process.wait()
         stderr_thread.join(timeout=2.0)
+        final_size = os.path.getsize(file_path) if os.path.exists(file_path) else 0
+        success = exit_code == 0 and final_size > 0
         with migration_lock:
             migration_state["dump_file_size"] = final_size / (1024 * 1024) # MB
             migration_state["end_time"] = time.time()
             migration_state["running"] = False
             migration_state["process"] = None
+            migration_state["dump_completed"] = success
+            if success:
+                # Write manifest file on successful dump
+                try:
+                    with open(manifest_path, 'w') as f:
+                        json.dump(manifest_data, f)
+                    log_message(f"Successfully saved progress manifest to {manifest_path}", "success")
+                except Exception as e:
+                    log_message(f"Could not save progress manifest: {e}", "error")
                 total_time = migration_state["end_time"] - migration_state["start_time"]
+                log_message(f"Dump completed successfully. Size: {migration_state['dump_file_size']:.2f} MB. Time: {total_time:.2f}s", "success")
             else:
+                log_message(f"Dump failed. Exit code: {exit_code}, File size: {final_size} bytes.", "error")
     except Exception as e:
+        log_message(f"An unexpected error occurred during dump: {e}", "error")
         with migration_lock:
             migration_state["running"] = False
             migration_state["process"] = None
 def run_restore(target_conn: str, file_path: str, options: dict):
     try:
         # Pre-restore step
         if options.get("timescaledb_pre_restore", True):
             pre_restore_cmd = ["psql", target_conn, "-c", "SELECT timescaledb_pre_restore();"]
+            pre_restore_proc = subprocess.run(pre_restore_cmd, capture_output=True, text=True)
             if pre_restore_proc.returncode != 0:
                 log_message(f"Pre-restore failed: {pre_restore_proc.stderr or pre_restore_proc.stdout}", "error")
                 return
         cmd = ["pg_restore", "-d", target_conn, "-v"]
         if options.get("no_owner", True): cmd.append("--no-owner")
         if options.get("clean", False): cmd.append("--clean")
         cmd.append(file_path)
+        log_message(f"Starting database restore...", "info")
+        process = subprocess.Popen(cmd, stderr=subprocess.PIPE, text=True, bufsize=1)
         with migration_lock:
             migration_state["process"] = process
+        stderr_thread = threading.Thread(target=process_stderr, args=(process, "restore"), daemon=True)
         stderr_thread.start()
         exit_code = process.wait()
         stderr_thread.join(timeout=2.0)
         post_restore_success = True
         if exit_code == 0 and options.get("timescaledb_post_restore", True):
             post_restore_cmd = ["psql", target_conn, "-c", "SELECT timescaledb_post_restore(); ANALYZE;"]
+            post_restore_proc = subprocess.run(post_restore_cmd, capture_output=True, text=True)
             if post_restore_proc.returncode != 0:
                 log_message(f"Post-restore failed: {post_restore_proc.stderr or post_restore_proc.stdout}", "error")
                 post_restore_success = False
             migration_state["end_time"] = time.time()
             migration_state["running"] = False
             migration_state["process"] = None
+            migration_state["restore_completed"] = (exit_code == 0 and post_restore_success)
+            if migration_state["restore_completed"]:
                 total_time = migration_state["end_time"] - migration_state["start_time"]
+                log_message(f"Restore completed successfully. Time: {total_time:.2f}s", "success")
             else:
+                log_message(f"Restore failed. Exit code: {exit_code}, Post-restore success: {post_restore_success}", "error")
     except Exception as e:
+        log_message(f"An unexpected error occurred during restore: {e}", "error")
         with migration_lock:
             migration_state["running"] = False
             migration_state["process"] = None
 def stop_current_process():
     with migration_lock:
         process = migration_state.get("process")
         if not process or not migration_state.get("running"):
             return False
         try:
             pid = process.pid
+            log_message(f"Attempting to stop {migration_state['operation']} process (PID: {pid})...", "warning")
+            if hasattr(os, 'killpg') and hasattr(os, 'getpgid'):
+                os.killpg(os.getpgid(pid), signal.SIGKILL)
             else:
+                process.kill()
+            process.wait(timeout=5)
+            log_message("Process stopped by user.", "warning")
             return True
         except Exception as e:
             log_message(f"Error stopping process: {e}", "error")
             return False
         finally:
             migration_state["running"] = False
             migration_state["process"] = None
             migration_state["end_time"] = time.time()
 # --- API Endpoints ---
 @app.get("/", response_class=HTMLResponse)
 async def home(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})
+def reset_state(operation: str):
+    """Resets the global state for a new operation."""
+    with migration_lock:
+        migration_state.update({
+            "id": str(uuid.uuid4()), "running": False, "operation": operation,
+            "start_time": None, "end_time": None, "dump_file": None, "dump_file_size": 0,
+            "dump_completed": False, "restore_completed": False, "last_activity": time.time(),
+            "log": [], "process": None,
+            "progress": {
+                "current_table": None, "total_chunks": 0, "chunks_processed": 0,
+                "total_size_bytes": 0, "size_processed_bytes": 0, "percent_complete_by_size": 0,
+                "percent_complete_by_count": 0, "manifest_loaded": False, "manifest": None,
+                "tables_completed": 0,
+            }
+        })
+def redact_password(conn_str: str) -> str:
+    return re.sub(r':([^@]+)@', r':********@', conn_str)
 @app.post("/start-dump")
 async def start_dump_endpoint(data: Dict[str, Any], background_tasks: BackgroundTasks):
     source_conn = data.get("source_conn")
     options = data.get("options", {})
+    if not source_conn or not test_connection_logic(source_conn):
+        raise HTTPException(status_code=400, detail="Valid source connection string is required.")
     if migration_state["running"]:
         stop_current_process()
+        time.sleep(1)
+    reset_state("dump")
+    manifest_data = get_chunk_manifest(source_conn)
+    filename = "".join(c for c in options.get("filename", "timescale_backup") if c.isalnum() or c in ('_','-'))
+    ext = {"p": ".sql", "d": "", "t": ".tar"}.get(options.get("format", "c"), ".dump")
+    dump_path = Path("dumps").resolve() / f"{filename}{ext}"
+    manifest_path = Path("dumps").resolve() / f"{filename}{ext}.manifest.json"
     with migration_lock:
+        migration_state["dump_file"] = str(dump_path)
+        migration_state["progress"]["manifest"] = manifest_data
+        migration_state["progress"]["total_chunks"] = manifest_data["total_chunks"]
+        migration_state["progress"]["total_size_bytes"] = manifest_data["total_size_bytes"]
+        migration_state["progress"]["manifest_loaded"] = manifest_data["total_chunks"] > 0
+        migration_state["start_time"] = time.time()
+        migration_state["running"] = True
+    background_tasks.add_task(run_dump, source_conn, str(dump_path), str(manifest_path), manifest_data, options)
+    return JSONResponse({"success": True, "message": "Dump process initiated."})
 @app.post("/start-restore")
 async def start_restore_endpoint(data: Dict[str, Any], background_tasks: BackgroundTasks):
     target_conn = data.get("target_conn")
     dump_file = data.get("dump_file")
     options = data.get("options", {})
+    if not target_conn or not dump_file or not test_connection_logic(target_conn):
+        raise HTTPException(status_code=400, detail="Valid target connection and dump file are required.")
+    dump_path = Path(dump_file).resolve()
+    if not dump_path.exists() or not str(dump_path).startswith(str(Path("dumps").resolve())):
+        raise HTTPException(status_code=400, detail="Invalid or non-existent dump file.")
     if migration_state["running"]:
         stop_current_process()
         time.sleep(1)
+    reset_state("restore")
+    manifest_path = Path(f"{dump_file}.manifest.json")
+    manifest_data = None
+    if manifest_path.exists():
+        try:
+            with open(manifest_path, 'r') as f:
+                manifest_data = json.load(f)
+            log_message("Successfully loaded progress manifest for restore.", "success")
+        except Exception as e:
+            log_message(f"Could not load manifest file: {e}", "error")
     with migration_lock:
+        if manifest_data:
+            migration_state["progress"]["manifest"] = manifest_data
+            migration_state["progress"]["total_chunks"] = manifest_data["total_chunks"]
+            migration_state["progress"]["total_size_bytes"] = manifest_data["total_size_bytes"]
+            migration_state["progress"]["manifest_loaded"] = True
+        migration_state["start_time"] = time.time()
+        migration_state["running"] = True
+    background_tasks.add_task(run_restore, target_conn, str(dump_path), options)
+    return JSONResponse({"success": True, "message": "Restore process initiated."})
 @app.post("/stop-process")
 async def stop_process_endpoint():
+    if stop_current_process():
+        return JSONResponse({"success": True, "message": "Process stop initiated."})
+    return JSONResponse({"success": False, "message": "No process was running to stop."})
 @app.get("/status")
 async def get_status():
     with migration_lock:
         state_copy = {k: v for k, v in migration_state.items() if k != "process"}
         if "progress" in state_copy and "manifest" in state_copy["progress"]:
+            state_copy["progress"]["manifest"] = None # Avoid sending large object
     return JSONResponse(content=state_copy)
 @app.get("/list-dumps")
 async def list_dumps():
+    dumps_dir = Path("dumps")
+    dump_files = []
+    for f in dumps_dir.iterdir():
+        if f.is_file() and not f.name.endswith('.manifest.json'):
+            try:
+                stat = f.stat()
+                dump_files.append({
+                    "name": f.name, "path": str(f), "size_mb": stat.st_size / (1024*1024),
+                    "date": datetime.datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %H:%M:%S")
+                })
+            except OSError as e:
+                logger.error(f"Could not stat file {f.name}: {e}")
+    dump_files.sort(key=lambda x: x["date"], reverse=True)
+    return JSONResponse({"success": True, "dumps": dump_files})
 @app.get("/downloads/{file_name:path}")
 async def download_file(file_name: str):
+    file_path = (Path("dumps").resolve() / file_name).resolve()
+    if not file_path.exists() or not str(file_path).startswith(str(Path("dumps").resolve())):
+        raise HTTPException(status_code=404, detail="File not found")
+    return FileResponse(path=str(file_path), filename=file_name)
 @app.post("/test-connection")
 async def test_connection_endpoint(data: Dict[str, str]):
     connection_string = data.get("connection_string")
+    if not connection_string or not test_connection_logic(connection_string):
+        raise HTTPException(status_code=400, detail="Failed to connect to database")
     try:
+        with psycopg2.connect(connection_string) as conn:
+            with conn.cursor() as cur:
+                cur.execute("SELECT version()")
+                version = cur.fetchone()[0]
+                cur.execute("SELECT current_database()")
+                database = cur.fetchone()[0]
+                is_timescaledb, ts_version = False, None
+                cur.execute("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'timescaledb');")
+                if cur.fetchone()[0]:
+                    is_timescaledb = True
+                    cur.execute("SELECT extversion FROM pg_extension WHERE extname = 'timescaledb'")
+                    ts_version = cur.fetchone()[0]
+        return JSONResponse({
             "success": True, "version": version, "is_timescaledb": is_timescaledb,
             "timescaledb_version": ts_version, "database": database
         })
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Database error: {e}")
+@app.post("/database-info")
+async def get_database_info(data: Dict[str, str]):
+    connection_string = data.get("connection_string")
+    if not connection_string or not test_connection_logic(connection_string):
+        raise HTTPException(status_code=400, detail="Connection failed")
+    try:
+        with psycopg2.connect(connection_string) as conn:
+            with conn.cursor() as cur:
+                cur.execute("""
+                    SELECT count(*) FROM information_schema.tables
+                    WHERE table_schema NOT IN ('pg_catalog', 'information_schema', '_timescaledb_internal', '_timescaledb_catalog', '_timescaledb_cache')
+                    AND table_type = 'BASE TABLE'
+                """)
+                table_count = cur.fetchone()[0]
+                cur.execute("SELECT pg_size_pretty(pg_database_size(current_database()))")
+                db_size = cur.fetchone()[0]
+        return JSONResponse({"success": True, "table_count": table_count, "database_size": db_size})
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Database query error: {e}")
 if __name__ == "__main__":
     import uvicorn