aaPanel/mod/project/project_import/database/importer.py

# coding: utf-8

import gzip
import os
import shutil
import subprocess
import tempfile

from ..core.exceptions import ProjectImportError
from ..core.security import ensure_child_path
from .config_updater import update_database_config


class DatabaseImporter:
    def __init__(self, config, session, project_result, project_path, progress, cancelled):
        self.config = config or {}
        self.session = session
        self.project_result = project_result or {}
        self.project_path = project_path
        self.progress = progress
        self.cancelled = cancelled

    def run(self):
        if not self.config.get("enabled"):
            return {"enabled": False}
        import public
        from database_v2 import database as DatabaseModel

        name = str(self.config.get("database_name", "")).strip().lower()
        user = str(self.config.get("database_user", name)).strip().lower()
        password = str(self.config.get("database_password", ""))
        charset = str(self.config.get("charset", "utf8mb4"))
        if not name or not user or not password:
            raise ProjectImportError("Database name, user and password are required", "DATABASE_CONFIG_REQUIRED")
        # 在建库前先定位SQL文件 文件缺失直接失败
        sql_path = self._resolve_sql_file()
        args = public.to_dict_obj({
            "name": name,
            "db_user": user,
            "codeing": charset,
            "password": password,
            "sid": 0,
            "active": True,
            "address": "127.0.0.1",
            "ps": str(self.project_result.get("project_name", "Imported project")),
            "dtype": "MySQL",
            "pid": int(self.project_result.get("site_id", 0) or 0),
        })
        result = DatabaseModel().AddDatabase(args)
        status = result.get("status") if isinstance(result, dict) else None
        # 注意：success 时 status 为 0（return_message(0, ...)），不能写成
        # `status in (-1, False)`——Python 中 0 == False，会把成功误判为失败。
        if not isinstance(result, dict) or status is False or status == -1:
            message = result.get("message", result) if isinstance(result, dict) else result
            if isinstance(message, dict) and "result" in message:
                message = message["result"]
            raise ProjectImportError("Failed to create database: {}".format(message), "DATABASE_CREATE_FAILED")
        if sql_path:
            self._stream_sql(sql_path, name)
        database_id = public.M("databases").where("name=?", (name,)).getField("id") or 0
        updated = []
        if self.config.get("update_project_config"):
            updated = update_database_config(
                self.project_path,
                str(self.project_result.get("project_type", "")),
                {"name": name, "user": user, "password": password, "host": "127.0.0.1"},
            )
        return {
            "enabled": True,
            "database_id": int(database_id),
            "database_name": name,
            "database_user": user,
            "config_updated": [os.path.relpath(item, self.project_path).replace("\\", "/") for item in updated],
        }

    def _resolve_sql_file(self):
        """根据 sql_file_path 定位要导入的 SQL/SQL.GZ 文件。

        扫描到的文件（analysis.sql_files[].path）和用户上传到服务器指定目录的
        文件都以实际服务端路径传入；相对路径则相对项目根目录解析（限制在根目录内）。
        """
        configured = str(self.config.get("sql_file_path", "")).strip()
        if not configured:
            return ""
        if not os.path.isabs(configured):
            root = self.session.get("internal", {}).get("project_root", "")
            if not root:
                return ""
            configured = ensure_child_path(root, os.path.join(root, configured))
        path = os.path.realpath(configured)
        if os.path.isfile(path):
            return path
        raise ProjectImportError(
            "The selected SQL file does not exist on the server: {}".format(configured),
            "SQL_FILE_NOT_FOUND",
        )

    def _stream_sql(self, sql_path, database_name):
        import public
        mysql_bin = public.get_mysql_bin()
        if not mysql_bin or not os.path.isfile(mysql_bin):
            mysql_bin = shutil.which("mysql")
        if not mysql_bin:
            raise ProjectImportError("MySQL client was not found", "MYSQL_CLIENT_NOT_FOUND")
        root_password = public.M("config").where("id=?", (1,)).getField("mysql_root") or ""
        fd, defaults_file = tempfile.mkstemp(prefix="project_import_mysql_", suffix=".cnf")
        os.close(fd)
        try:
            with open(defaults_file, "w", encoding="utf-8") as handle:
                handle.write("[client]\nuser=root\npassword={}\ndefault-character-set=utf8mb4\n".format(root_password))
            os.chmod(defaults_file, 0o600)
            command = [mysql_bin, "--defaults-extra-file={}".format(defaults_file), database_name]
            process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
            total = max(1, os.path.getsize(sql_path))
            source = gzip.open(sql_path, "rb") if sql_path.lower().endswith(".gz") else open(sql_path, "rb")
            try:
                while True:
                    self.cancelled()
                    chunk = source.read(1024 * 1024)
                    if not chunk:
                        break
                    process.stdin.write(chunk)
                    position = source.fileobj.tell() if hasattr(source, "fileobj") else source.tell()
                    self.progress(min(0.99, position / total), "Importing database")
                process.stdin.close()
                stderr = process.stderr.read().decode("utf-8", errors="ignore") if process.stderr else ""
                code = process.wait()
                if code != 0:
                    raise ProjectImportError("Database import failed: {}".format(stderr[-2000:]), "DATABASE_IMPORT_FAILED")
                self.progress(1, "Database import completed")
            finally:
                source.close()
                if process.poll() is None:
                    process.terminate()
        finally:
            try:
                os.remove(defaults_file)
            except OSError:
                pass