mcp-server/spark_executor/tools/connections.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from spark_executor.core.connection_store import store
from spark_executor.models import Connection


_UNSET = object()


def save_connection(
    *,
    name: str,
    master: str,
    deploy_mode: str = _UNSET,  # type: ignore[assignment]
    yarn_rm_url: str | None = _UNSET,  # type: ignore[assignment]
    spark_conf: dict[str, str] | None = _UNSET,  # type: ignore[assignment]
    ssl_verify: bool | None = _UNSET,  # type: ignore[assignment]
    ssl_ca_bundle: str | None = _UNSET,  # type: ignore[assignment]
    auth_type: str = _UNSET,  # type: ignore[assignment]
    auth_user: str | None = _UNSET,  # type: ignore[assignment]
    auth_password: str | None = _UNSET,  # type: ignore[assignment]
    auth_principal: str | None = _UNSET,  # type: ignore[assignment]
    auth_keytab: str | None = _UNSET,  # type: ignore[assignment]
    url_allowlist: list[str] | None = _UNSET,  # type: ignore[assignment]
) -> dict[str, object]:
    logger.debug(
        f"save_connection enter name={name} master={master} "
        f"spark_conf_keys={list((spark_conf if isinstance(spark_conf, dict) else {}).keys())}"
    )
    existing = store.get(name)
    if existing is not None:
        fields: dict[str, object] = {"master": master}
        if deploy_mode is not _UNSET:
            fields["deploy_mode"] = deploy_mode
        if yarn_rm_url is not _UNSET:
            fields["yarn_rm_url"] = yarn_rm_url
        if spark_conf is not _UNSET:
            fields["spark_conf"] = spark_conf or {}
        if ssl_verify is not _UNSET:
            fields["ssl_verify"] = ssl_verify
        if ssl_ca_bundle is not _UNSET:
            fields["ssl_ca_bundle"] = ssl_ca_bundle
        if auth_type is not _UNSET:
            fields["auth_type"] = auth_type
        if auth_user is not _UNSET:
            fields["auth_user"] = auth_user
        if auth_password is not _UNSET:
            fields["auth_password"] = auth_password
        if auth_principal is not _UNSET:
            fields["auth_principal"] = auth_principal
        if auth_keytab is not _UNSET:
            fields["auth_keytab"] = auth_keytab
        if url_allowlist is not _UNSET:
            fields["url_allowlist"] = url_allowlist or []
        return update_connection(name, **fields)
    new_fields: dict[str, object] = {
        "name": name,
        "master": master,
        "deploy_mode": deploy_mode if deploy_mode is not _UNSET else "cluster",
        "yarn_rm_url": yarn_rm_url if yarn_rm_url is not _UNSET else None,
        "spark_conf": (spark_conf if spark_conf is not _UNSET else None) or {},
        "ssl_verify": ssl_verify if ssl_verify is not _UNSET else None,
        "ssl_ca_bundle": ssl_ca_bundle if ssl_ca_bundle is not _UNSET else None,
        "auth_type": auth_type if auth_type is not _UNSET else "none",
        "auth_user": auth_user if auth_user is not _UNSET else None,
        "auth_password": auth_password if auth_password is not _UNSET else None,
        "auth_principal": auth_principal if auth_principal is not _UNSET else None,
        "auth_keytab": auth_keytab if auth_keytab is not _UNSET else None,
        "url_allowlist": (url_allowlist if url_allowlist is not _UNSET else None) or [],
    }
    conn = Connection(**new_fields)
    store.save(conn)
    return {"name": name, "status": "SAVED"}


def update_connection(name: str, **fields) -> dict[str, object]:
    """Update an existing Connection's mutable fields.

    `name` is the identifier (immutable). PATCH semantics: only the fields
    you pass are changed. To clear an optional field (e.g. `yarn_rm_url`),
    use `delete_connection(name=...)` followed by `save_connection(...)`.

    Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf,
    ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password,
    auth_principal, auth_keytab, url_allowlist.
    """
    logger.debug(f"update_connection enter name={name} fields={sorted(fields.keys())}")
    return store.update(name, **fields).model_dump()


def list_connections() -> list[dict[str, object]]:
    logger.debug("list_connections enter")
    return [c.model_dump() for c in store.list_all()]


def get_connection(name: str) -> dict[str, object]:
    logger.debug(f"get_connection enter name={name}")
    conn = store.get(name)
    if conn is None:
        raise KeyError(f"Unknown connection: {name}")
    return conn.model_dump()


def delete_connection(name: str) -> dict[str, str]:
    logger.debug(f"delete_connection enter name={name}")
    removed = store.delete(name)
    if not removed:
        raise KeyError(f"Unknown connection: {name}")
    return {"name": name, "status": "DELETED"}