diff --git a/airflow-core/src/airflow/models/variable.py b/airflow-core/src/airflow/models/variable.py index b06e73cd5f50a..2f5d4967f8135 100644 --- a/airflow-core/src/airflow/models/variable.py +++ b/airflow-core/src/airflow/models/variable.py @@ -160,6 +160,9 @@ def get( :param deserialize_json: Deserialize the value to a Python dict :param team_name: Team name associated to the task trying to access the variable (if any) """ + if not key or not isinstance(key, str): + raise ValueError("Variable key must be a non-empty string") + # 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 @@ -220,6 +223,9 @@ def set( :param team_name: Team name associated to the variable (if any) :param session: optional session, use if provided or create a new one """ + if not key or not isinstance(key, str): + raise ValueError("Variable key must be a non-empty string") + # 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 diff --git a/airflow-core/tests/unit/models/test_variable.py b/airflow-core/tests/unit/models/test_variable.py index e56ee51d734f7..76180e68ea291 100644 --- a/airflow-core/tests/unit/models/test_variable.py +++ b/airflow-core/tests/unit/models/test_variable.py @@ -339,6 +339,20 @@ def test_variable_delete_with_team(self, testing_team, session): Variable.set(key=key, value=value, team_name=testing_team.name, session=session) assert Variable.delete(key=key, session=session) == 0 + @pytest.mark.parametrize("bad_key", [None, "", 123]) + def test_variable_set_rejects_invalid_key(self, bad_key, session): + with pytest.raises(ValueError, match="Variable key must be a non-empty string"): + Variable.set(key=bad_key, value="value", session=session) + + @pytest.mark.parametrize("bad_key", [None, "", 123]) + def test_variable_get_rejects_invalid_key(self, bad_key): + with pytest.raises(ValueError, match="Variable key must be a non-empty string"): + Variable.get(bad_key) + + def test_variable_setdefault_rejects_invalid_key(self, session): + with pytest.raises(ValueError, match="Variable key must be a non-empty string"): + Variable.setdefault(key=None, default="value") + def test_masking_from_db(self, session): """Test secrets are masked when loaded directly from the DB""" # Normally people will use `Variable.get`, but just in case, catch direct DB access too