Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import json
from collections.abc import Iterable

from pydantic import Field, JsonValue, model_validator
from pydantic import Field, JsonValue, field_validator, model_validator

from airflow._shared.secrets_masker import redact
from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel, make_partial_model
Expand Down Expand Up @@ -61,6 +61,13 @@ class VariableBody(StrictBaseModel):
description: str | None = Field(default=None)
team_name: str | None = Field(max_length=50, default=None)

@field_validator("value")
@classmethod
def serialize_non_string_value(cls, value: JsonValue) -> JsonValue:
if isinstance(value, str):
return value
return json.dumps(value, indent=2)

@model_validator(mode="after")
def validate_team_name(self) -> VariableBody:
if self.team_name is not None and not conf.getboolean("core", "multi_team"):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -126,13 +126,11 @@ def handle_bulk_create(self, action: BulkCreateAction, results: BulkActionRespon

for variable in action.entities:
if variable.key in create_keys:
should_serialize_json = isinstance(variable.value, (dict, list))
Variable.set(
key=variable.key,
value=variable.value,
description=variable.description,
session=self.session,
serialize_json=should_serialize_json,
)
results.success.append(variable.key)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1528,3 +1528,74 @@ def test_bulk_delete_resolves_existence_in_single_query(self, session, num_varia

session.commit()
assert session.scalars(select(Variable.key).where(Variable.key.in_(keys))).all() == []


class TestVariableJsonValueSerialization(TestVariableEndpoint):
"""``value`` accepts any JSON type, so every write path must store it as JSON, not a Python repr."""

JSON_VALUES = [
(["a", "b"], '[\n "a",\n "b"\n]'),
({"k": 1}, '{\n "k": 1\n}'),
(True, "true"),
(7, "7"),
(None, "null"),
]
IDS = ["list", "dict", "bool", "int", "null"]

@pytest.mark.parametrize(("value", "expected_stored"), JSON_VALUES, ids=IDS)
def test_post_stores_json(self, test_client, session, value, expected_stored):
response = test_client.post("/variables", json={"key": "json_var", "value": value})
assert response.status_code == 201

stored = session.scalar(select(Variable).where(Variable.key == "json_var"))
assert stored.val == expected_stored
assert json.loads(stored.val) == value

@pytest.mark.parametrize(("value", "expected_stored"), JSON_VALUES, ids=IDS)
def test_patch_stores_json(self, test_client, session, value, expected_stored):
Variable.set(key="json_var", value="original", session=session)
session.commit()

response = test_client.patch("/variables/json_var", json={"key": "json_var", "value": value})
assert response.status_code == 200

session.expire_all()
stored = session.scalar(select(Variable).where(Variable.key == "json_var"))
assert stored.val == expected_stored
assert json.loads(stored.val) == value

@pytest.mark.parametrize(("value", "expected_stored"), JSON_VALUES, ids=IDS)
def test_bulk_create_stores_json(self, test_client, session, value, expected_stored):
response = test_client.patch(
"/variables",
json={"actions": [{"action": "create", "entities": [{"key": "json_var", "value": value}]}]},
)
assert response.status_code == 200
assert response.json()["create"]["success"] == ["json_var"]

stored = session.scalar(select(Variable).where(Variable.key == "json_var"))
assert stored.val == expected_stored
assert json.loads(stored.val) == value

@pytest.mark.parametrize("verb", ["post", "patch", "bulk"])
def test_string_value_is_stored_verbatim(self, test_client, session, verb):
"""A string is already the stored form -- it must not be re-encoded into a quoted JSON string."""
if verb == "post":
test_client.post("/variables", json={"key": "str_var", "value": "plain string"})
elif verb == "patch":
Variable.set(key="str_var", value="original", session=session)
session.commit()
test_client.patch("/variables/str_var", json={"key": "str_var", "value": "plain string"})
else:
test_client.patch(
"/variables",
json={
"actions": [
{"action": "create", "entities": [{"key": "str_var", "value": "plain string"}]}
]
},
)

session.expire_all()
stored = session.scalar(select(Variable).where(Variable.key == "str_var"))
assert stored.val == "plain string"
Loading