# 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"}