vatsrahul1001 commented on code in PR #73897: URL: https://github.com/apache/airflow/pull/73897#discussion_r4134897879
########## providers/common/ai/src/airflow/providers/common/ai/utils/masking.py: ########## @@ -0,0 +1,120 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Apply Airflow's secret masker to what a tool hands back to a model.""" + +from __future__ import annotations + +import dataclasses +import json +from typing import Any, overload + +from pydantic import BaseModel +from pydantic_ai.messages import MULTI_MODAL_CONTENT_TYPES +from pydantic_core import to_jsonable_python + +from airflow.providers.common.compat.sdk import redact + +# Stands in for a container that contains itself, which would otherwise recurse forever. +_CYCLE = "<circular reference>" + + +@overload +def mask_secrets(value: str) -> str: ... + + +@overload +def mask_secrets(value: Any) -> Any: ... + + +def mask_secrets(value: Any) -> Any: + """ + Return ``value`` with every secret Airflow has registered replaced by ``***``. + + Strings are masked wherever they sit in nested dicts, lists, tuples, sets and dataclasses, + dict keys included, and the shape and types are kept. Bytes are masked as UTF-8 text. A + Pydantic model is turned into the JSON-compatible data the model would be shown. Images, + documents and other multimodal content, and any other object, pass through as they are. + Registered secrets are the ones Airflow knows about, such as connection passwords and + sensitive connection extras; a credential that only appears in the data itself is not + recognized. + + ``redact()`` does part of this, but stops descending at a fixed depth, and it hides every + string under a key that looks sensitive: a model reading a query result needs + ``{"auth_type": "oauth"}`` as it is. Two dict keys that both mask to ``***`` collapse into + one. + """ + return _mask(value, frozenset()) + + +def dumps_masked(value: Any, **kwargs: Any) -> str: + """ + Serialize ``value`` to JSON for a model, with registered secrets masked first. + + Masking a JSON string afterwards is not enough: JSON escapes quotes, backslashes, + control characters and, by default, non-ASCII characters, so a password containing + any of them no longer matches the registered value once it is inside the document. + Bytes become their UTF-8 text and dataclasses their fields; any other object JSON + cannot represent is rendered with ``str()``, after the values inside it are masked. + + :param kwargs: Passed to :func:`json.dumps`. + """ + return json.dumps(mask_secrets(value), default=_json_default, **kwargs) + + +def _json_default(value: object) -> Any: + if isinstance(value, bytes): + return value.decode("utf-8", "replace") + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return {field.name: getattr(value, field.name) for field in dataclasses.fields(value)} + return mask_secrets(str(value)) + + +def _mask(value: Any, seen: frozenset[int]) -> Any: Review Comment: can this hit a RecursionError on a deeply nested result? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
