Repository: apache/airflow
File: airflow-core/src/airflow/models/variable.py (Python)

Code before:
        # means SQLA etc is loaded, but we can't avoid that unless/until we add import shims as a big
        # back-compat layer

        # If this is set it means we are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
        # and should use the Task SDK API server path
        if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
            warnings.warn(
                "Using Variable.get from `airflow.models` is deprecated."
                "Please use `get` on Variable from sdk(`airflow.sdk.Variable`) instead",
                DeprecationWarning,
                stacklevel=1,
            )
            from airflow.sdk import Variable as TaskSDKVariable

            default_kwargs = {} if default_var is cls.__NO_DEFAULT_SENTINEL else {"default": default_var}
            var_val = TaskSDKVariable.get(key, deserialize_json=deserialize_json, **default_kwargs)
            if isinstance(var_val, str):
                mask_secret(var_val, key)

            return var_val

        var_val = Variable.get_variable_from_secrets(key=key)
        if var_val is None:
            if default_var is not cls.__NO_DEFAULT_SENTINEL:
                return default_var
            raise KeyError(f"Variable {key} does not exist. Try using a full path to your secret if you have one.")
        if deserialize_json:
            obj = json.loads(var_val)
            mask_secret(obj, key)
            return obj
        mask_secret(var_val, key)
        return var_val

    @staticmethod
    def set(
        key: str,
        value: Any,
        description: str | None = None,
        serialize_json: bool = False,
        team_name: str | None = None,
        session: Session | None = None,
    ) -> None:
        """
        Set a value for an Airflow Variable with a given Key.

        This operation overwrites an existing variable using the session's dialect-specific upsert operation.

        :param key: Variable Key
        :param value: Value to set for the Variable
        :param description: Description of the Variable
        :param serialize_json: Serialize the value to a JSON string

Review comment (security):
This applies to all variables, regardless of how they're fetched (secrets backend/DB/etc.), would remove.

Code after:
        # TODO: This is not the best way of having compat, but it's "better than erroring" for now. This still
        # means SQLA etc is loaded, but we can't avoid that unless/until we add import shims as a big
        # back-compat layer

        # If this is set it means we are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
        # and should use the Task SDK API server path
        if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
            warnings.warn(
                "Using Variable.get from `airflow.models` is deprecated."
                "Please use `get` on Variable from sdk(`airflow.sdk.Variable`) instead",
                DeprecationWarning,
                stacklevel=1,
            )
            from airflow.sdk import Variable as TaskSDKVariable

            default_kwargs = {} if default_var is cls.__NO_DEFAULT_SENTINEL else {"default": default_var}
            var_val = TaskSDKVariable.get(key, deserialize_json=deserialize_json, **default_kwargs)
            if isinstance(var_val, str):
                mask_secret(var_val, key)

            return var_val

        if team_name and not conf.getboolean("core", "multi_team"):
            raise ValueError(
                "Multi-team mode is not configured in the Airflow environment but the task trying to access the variable belongs to a team"
            )

        var_val = Variable.get_variable_from_secrets(key=key, team_name=team_name)
        if var_val is None:
            if default_var is not cls.__NO_DEFAULT_SENTINEL:
                return default_var
            raise KeyError(f"Variable {key} does not exist.")
        if deserialize_json:
            obj = json.loads(var_val)
            mask_secret(obj, key)
            return obj
        mask_secret(var_val, key)
        return var_val

    @staticmethod
    def set(
        key: str,
        value: Any,
        description: str | None = None,
        serialize_json: bool = False,
        team_name: str | None = None,
        session: Session | None = None,
    ) -> None:
        """
        Set a value for an Airflow Variable with a given Key.