db_import_server.py

#!/usr/bin/env python3
"""
Database Import Studio — Local HTTP Server
Uses Python 3 standard library only. No pip install needed.
Serves the HTML UI and handles import API calls.
"""
import http.server, json, os, re, subprocess, sys, threading, time, webbrowser
from pathlib import Path
from urllib.parse import urlparse, parse_qs

# ── Paths ─────────────────────────────────────────────────────────────────────
if getattr(sys, 'frozen', False):
    BASE_DIR = Path(sys._MEIPASS)
else:
    BASE_DIR = Path(__file__).parent.resolve()
if os.name == "nt":
    LOG_DIR = Path(os.environ.get("APPDATA", Path.home())) / "DatabaseImportStudio" / "logs"
else:
    LOG_DIR = Path.home() / ".local" / "share" / "DatabaseImportStudio" / "logs"
LOG_DIR.mkdir(parents=True, exist_ok=True)

# ── Global state ──────────────────────────────────────────────────────────────
_lock   = threading.Lock()
_state  = {"status": "idle", "percent": 0, "bytesSent": 0, "fileSize": 0,
           "message": "Ready", "logFile": "", "expectedTables": 0, "ts": 0.0}
_proc   = None          # running mysql subprocess
_thread = None          # import thread
_cancel = threading.Event()
_server = None          # HTTPServer instance (set at startup)

# ── MySQL detection ───────────────────────────────────────────────────────────
def find_mysql():
    candidates = []
    if os.name == "nt":
        for root in ["C:/wamp64", "C:/wamp", "C:/xampp/mysql", "C:/mysql",
                     "C:/Program Files/MySQL/MySQL Server 8.0"]:
            rp = Path(root)
            if rp.exists():
                for p in rp.rglob("mysql.exe"):
                    candidates.append(str(p))
                for p in rp.rglob("mariadb.exe"):
                    candidates.append(str(p))
        candidates += ["mysql.exe", "mariadb.exe"]
    else:
        candidates = ["/usr/bin/mysql", "/usr/local/bin/mysql",
                      "/usr/bin/mariadb", "/usr/local/bin/mariadb", "mysql", "mariadb"]
    for c in candidates:
        try:
            r = subprocess.run([c, "--version"], capture_output=True, timeout=3)
            if r.returncode == 0:
                return c
        except Exception:
            pass
    return ""

# ── Table count scanner (runs in background) ──────────────────────────────────
def scan_tables_async(sql_file, file_size, log_fn):
    try:
        tables = set()
        pattern = re.compile(rb"(?i)CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?`?([a-zA-Z0-9_\$]+)`?")
        with open(sql_file, "rb") as f:
            leftover = b""
            while not _cancel.is_set():
                chunk = f.read(2 * 1024 * 1024) # 2MB chunks
                if not chunk:
                    break
                buffer = leftover + chunk
                matches = pattern.findall(buffer)
                for m in matches:
                    tables.add(m.decode("utf-8", errors="replace"))
                leftover = buffer[-512:]
                with _lock:
                    _state["expectedTables"] = len(tables)

        log_fn(f"Table scanner finished: found {len(tables)} tables.")
    except Exception as exc:
        log_fn(f"Table scanner notice: {exc}")

# ── Import worker (runs in background thread) ─────────────────────────────────
def run_import(cfg):
    global _proc
    _cancel.clear()
    log_file = str(LOG_DIR / f"import-{time.strftime('%Y%m%d-%H%M%S')}.log")

    def log(msg):
        line = f"[{time.strftime('%H:%M:%S')}] {msg}"
        with open(log_file, "a", encoding="utf-8") as f:
            f.write(line + "\n")

    def upd(status, pct, sent, msg):
        with _lock:
            _state.update(status=status, percent=pct, bytesSent=sent,
                          message=msg, ts=time.time(), logFile=log_file)

    try:
        mysql_exe = cfg["mysqlExe"]
        sql_file  = cfg["sqlFile"]
        database  = cfg["database"]
        user      = cfg.get("user", "root")
        password  = cfg.get("password", "")
        host      = cfg.get("host", "127.0.0.1")
        port      = cfg.get("port", "3306")

        base_args = [
            mysql_exe,
            f"--host={host}",
            f"--port={port}",
            f"--user={user}",
            "--default-character-set=utf8mb4",
            "--connect-timeout=15",
        ]
        if password:
            base_args.append(f"--password={password}")

        # ── Test connection ──────────────────────────────────────────────────
        upd("starting", 0, 0, "Testing connection…")
        r = subprocess.run(base_args + ["--execute=SELECT 1;"],
                           capture_output=True, timeout=20)
        if r.returncode != 0:
            err = r.stderr.decode("utf-8", errors="replace").strip()
            raise Exception(f"Cannot connect: {err or 'check host / user / password'}")
        log(f"Connected to MySQL at {host}:{port}")

        # ── Create / recreate database ───────────────────────────────────────
        db_sql = None
        if cfg.get("recreateDb"):
            db_sql = (f"DROP DATABASE IF EXISTS `{database}`; "
                      f"CREATE DATABASE `{database}` CHARACTER SET utf8mb4 "
                      f"COLLATE utf8mb4_unicode_ci;")
            log(f"Dropping and recreating database '{database}'…")
        elif cfg.get("createDb"):
            db_sql = (f"CREATE DATABASE IF NOT EXISTS `{database}` "
                      f"CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;")
            log(f"Creating database '{database}' if missing…")

        if db_sql:
            upd("starting", 0, 0, "Preparing database…")
            r = subprocess.run(base_args + [f"--execute={db_sql}"],
                               capture_output=True, timeout=30)
            if r.returncode != 0:
                err = r.stderr.decode("utf-8", errors="replace").strip()
                raise Exception(f"Failed to prepare database: {err}")

        # ── File info ────────────────────────────────────────────────────────
        file_size = os.path.getsize(sql_file)
        with _lock:
            _state["fileSize"] = file_size
            _state["expectedTables"] = 0
        log(f"SQL file: {sql_file}")
        log(f"File size: {file_size / 1048576:.1f} MB")

        # Start full table scan asynchronously
        threading.Thread(target=scan_tables_async, args=(sql_file, file_size, log), daemon=True).start()

        # ── Build import args ────────────────────────────────────────────────
        import_args = base_args[:]
        if cfg.get("force"):
            import_args.append("--force")
        import_args.append("--show-warnings")
        if cfg.get("speed"):
            import_args += ["--max_allowed_packet=1073741824",
                            "--net_buffer_length=1048576"]
        import_args.append(f"--database={database}")

        log("Starting import…")
        upd("running", 0, 0, "Starting import…")

        # ── Launch mysql process ─────────────────────────────────────────────
        _proc = subprocess.Popen(
            import_args,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        # CONTINUOUSLY drain stdout / stderr line-by-line / chunk-by-chunk so mysql never blocks on pipe buffers
        def _drain_continuous(stream):
            while True:
                chunk = stream.read(8192)
                if not chunk:
                    break
                with open(log_file, "ab") as f:
                    f.write(chunk)

        threading.Thread(target=_drain_continuous, args=(_proc.stdout,), daemon=True).start()
        threading.Thread(target=_drain_continuous, args=(_proc.stderr,), daemon=True).start()

        # ── Stream SQL file with byte progress ───────────────────────────────
        try:
            if cfg.get("disableFk"):
                header = (b"SET SESSION FOREIGN_KEY_CHECKS=0; "
                          b"SET SESSION UNIQUE_CHECKS=0; "
                          b"SET SESSION AUTOCOMMIT=0;\n")
                _proc.stdin.write(header)
                log("FK / unique checks disabled for this session.")

            sent = 0
            bad_mode = b"NO_AUTO_CREATE_USER"
            blank_mode = b" " * len(bad_mode)

            with open(sql_file, "rb") as f:
                while not _cancel.is_set():
                    chunk = f.read(1024 * 1024)  # 1 MB chunks
                    if not chunk:
                        break
                    # Auto-fix MySQL 8.0 compatibility
                    if bad_mode in chunk:
                        chunk = chunk.replace(bad_mode, blank_mode)

                    try:
                        _proc.stdin.write(chunk)
                    except (BrokenPipeError, OSError) as pipe_err:
                        log(f"Stream interrupted: mysql process closed stdin ({pipe_err}).")
                        break

                    sent += len(chunk)
                    pct = min(99, int(sent * 100 / file_size)) if file_size else 0
                    upd("running", pct, sent,
                        f"Importing… {sent/1048576:.1f} / {file_size/1048576:.1f} MB")

            if _cancel.is_set():
                try: _proc.kill()
                except Exception: pass
                upd("cancelled", 0, sent, "Import cancelled by user.")
                log("Import cancelled by user.")
                return

            if cfg.get("disableFk") and _proc.poll() is None:
                try:
                    _proc.stdin.write(
                        b"\nCOMMIT; "
                        b"SET SESSION FOREIGN_KEY_CHECKS=1; "
                        b"SET SESSION UNIQUE_CHECKS=1;\n"
                    )
                except Exception: pass

        finally:
            try: _proc.stdin.close()
            except Exception: pass

        _proc.wait()
        rc = _proc.returncode
        _proc = None

        if rc == 0:
            log("Import completed successfully.")
            upd("completed", 100, file_size, "Import completed successfully! ✓")
        else:
            log(f"mysql finished with exit code {rc}.")
            upd("completed_with_errors", 100, file_size,
                f"Completed with warnings/errors (exit code {rc}).")

    except Exception as exc:
        msg = str(exc)
        log(f"ERROR: {msg}")
        upd("failed", 0, 0, msg)
    finally:
        _proc = None


# ── HTTP handler ──────────────────────────────────────────────────────────────
class Handler(http.server.BaseHTTPRequestHandler):
    def log_message(self, *_): pass  # silence default access log

    def _cors(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")

    def send_json(self, obj, code=200):
        body = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self._cors()
        self.end_headers()
        self.wfile.write(body)

    def do_OPTIONS(self):
        self.send_response(200)
        self._cors()
        self.end_headers()

    def do_GET(self):
        parsed = urlparse(self.path)
        path   = parsed.path
        qs     = parse_qs(parsed.query)

        # ── Serve HTML UI ────────────────────────────────────────────────────
        if path in ("/", "/index.html"):
            html_file = BASE_DIR / "ui" / "index.html"
            body = html_file.read_bytes()
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return

        # ── Status ───────────────────────────────────────────────────────────
        if path == "/api/status":
            with _lock:
                s = dict(_state)
            self.send_json(s)
            return

        # ── Log content ──────────────────────────────────────────────────────
        if path == "/api/log":
            lf   = _state.get("logFile", "")
            text = ""
            if lf and os.path.exists(lf):
                with open(lf, encoding="utf-8", errors="replace") as f:
                    text = f.read()
            self.send_json({"log": text})
            return

        # ── Directory browser ─────────────────────────────────────────────────
        if path == "/api/browse":
            dir_path = qs.get("path", [str(Path.home())])[0].strip().strip('"').strip("'")
            try:
                # Handle Windows drive letters like D: or D:\
                if os.name == "nt" and re.match(r"^[a-zA-Z]:?$", dir_path):
                    dir_path = dir_path.rstrip(":") + ":\\"

                p = Path(dir_path)
                if not p.exists() or not p.is_dir():
                    if p.parent.exists() and p.parent.is_dir():
                        p = p.parent
                    else:
                        p = Path.home()

                entries = []
                for item in sorted(p.iterdir(),
                                   key=lambda x: (not x.is_dir(), x.name.lower())):
                    try:
                        entries.append({
                            "name":  item.name,
                            "isDir": item.is_dir(),
                            "path":  str(item),
                            "size":  item.stat().st_size if item.is_file() else 0,
                        })
                    except Exception:
                        pass

                drives = []
                if os.name == "nt":
                    import string
                    for letter in string.ascii_uppercase:
                        dp = f"{letter}:\\"
                        if os.path.exists(dp):
                            drives.append(dp)

                self.send_json({"path": str(p), "parent": str(p.parent), "entries": entries, "drives": drives})
            except Exception as exc:
                self.send_json({"error": str(exc)}, 500)
            return

        # ── Auto-detect mysql ─────────────────────────────────────────────────
        if path == "/api/detect-mysql":
            self.send_json({"exe": find_mysql()})
            return

        self.send_error(404, "Not found")

    def do_POST(self):
        global _thread
        parsed  = urlparse(self.path)
        path    = parsed.path
        length  = int(self.headers.get("Content-Length", 0))
        body    = json.loads(self.rfile.read(length)) if length else {}

        # ── Start import ──────────────────────────────────────────────────────
        if path == "/api/start":
            if _thread and _thread.is_alive():
                self.send_json({"error": "An import is already running."}, 409)
                return
            _thread = threading.Thread(target=run_import, args=(body,), daemon=True)
            _thread.start()
            self.send_json({"ok": True})
            return

        # ── Cancel ────────────────────────────────────────────────────────────
        if path == "/api/cancel":
            _cancel.set()
            if _proc:
                try: _proc.kill()
                except Exception: pass
            self.send_json({"ok": True})
            return

        # ── Stop server ───────────────────────────────────────────────────────
        if path == "/api/stop":
            self.send_json({"ok": True})
            if _server:
                threading.Thread(target=_server.shutdown, daemon=True).start()
            return

        self.send_error(404, "Not found")


# ── Entry point ───────────────────────────────────────────────────────────────
if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 7321
    _server = http.server.HTTPServer(("127.0.0.1", port), Handler)
    url = f"http://127.0.0.1:{port}"
    print(f"\n  Database Import Studio — UI ready")
    print(f"  Open in browser: {url}")
    print(f"  Press Ctrl+C to stop.\n")
    webbrowser.open(url)
    try:
        _server.serve_forever()
    except KeyboardInterrupt:
        print("\nServer stopped.")