Skip to content

Commit a20c177

Browse files
committed
Fix CI breakages from the dependency upgrade
- snowflake: snowflake-sqlalchemy (up to 1.11.1) cannot be imported on SQLAlchemy 2.1, which made ORMSelectCompileState private. Cap sqlalchemy<2.1 in the provider; the workspace lock resolves back to SQLAlchemy 2.0.54. Removing the cap is tracked in #75118. - airflow-core: type the dag_display_name and duration hybrid expressions against the class (type[...] -> ColumnElement) and fix _get_dialect's return type, so mypy holds on SQLAlchemy 2.0 and 2.1. - common.ai: declare the OpenAI Decisions and System One vendors added in pydantic-ai 2.53/2.55; reject sqlglot >= 30.22's exp.ListStage (Snowflake LIST/LS @stage) on the allow_writes path; adapt tests to pydantic-ai's streamed Anthropic requests and new request counting, and to sqlglot's MySQL "--+" tokenization. Generated-by: Claude Opus 5
1 parent e0934d2 commit a20c177

14 files changed

Lines changed: 153 additions & 126 deletions

File tree

‎airflow-core/src/airflow/models/backfill.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,7 @@ def duration(self) -> float | None:
220220
@duration.expression # type: ignore[no-redef]
221221
def duration(cls) -> Any:
222222
@provide_session
223-
def _get_dialect(*, session: Session = NEW_SESSION) -> str:
223+
def _get_dialect(*, session: Session = NEW_SESSION) -> str | None:
224224
return get_dialect_name(session=session)
225225

226226
dialect_name = _get_dialect()

‎airflow-core/src/airflow/models/dag.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@
8383

8484
from dateutil.relativedelta import relativedelta
8585
from sqlalchemy.orm.state import InstanceState
86+
from sqlalchemy.sql.elements import ColumnElement
8687

8788
from airflow.sdk import Context
8889
from airflow.serialization.definitions.assets import (
@@ -650,15 +651,15 @@ def dag_display_name(self) -> str:
650651
return self._dag_display_property_value or self.dag_id
651652

652653
@dag_display_name.expression # type: ignore[no-redef]
653-
def dag_display_name(self) -> str:
654+
def dag_display_name(cls: type[DagModel]) -> ColumnElement[Any]:
654655
"""
655656
Expression part of the ``dag_display`` name hybrid property.
656657
657658
:meta private:
658659
"""
659660
return case(
660-
(self._dag_display_property_value.is_not(None), self._dag_display_property_value),
661-
else_=self.dag_id,
661+
(cls._dag_display_property_value.is_not(None), cls._dag_display_property_value),
662+
else_=cls.dag_id,
662663
)
663664

664665
@classmethod

‎airflow-core/src/airflow/models/dagrun.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@
122122
from pydantic import NonNegativeInt
123123
from sqlalchemy.engine import ScalarResult
124124
from sqlalchemy.orm import Session
125-
from sqlalchemy.sql.elements import Case, ColumnElement
125+
from sqlalchemy.sql.elements import ColumnElement
126126
from sqlalchemy.sql.selectable import Select
127127

128128
from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
@@ -585,7 +585,7 @@ def duration(self) -> float | None:
585585

586586
@duration.expression # type: ignore[no-redef]
587587
@provide_session
588-
def duration(cls, *, session: Session = NEW_SESSION) -> Case:
588+
def duration(cls: type[DagRun], *, session: Session = NEW_SESSION) -> ColumnElement[Any]:
589589
dialect_name = get_dialect_name(session)
590590
if dialect_name == "mysql":
591591
return func.timestampdiff(text("SECOND"), cls.start_date, cls.end_date)

‎providers/common/ai/provider.yaml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,11 +197,13 @@ connection-types:
197197
- Nebius
198198
- Ollama
199199
- OpenAI
200+
- OpenAI Decisions
200201
- OpenRouter
201202
- OVHcloud
202203
- Pydantic AI Gateway
203204
- SambaNova
204205
- Snowflake Cortex
206+
- System One
205207
- Together AI
206208
- TypeSafe
207209
- Vercel AI Gateway

‎providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,11 +170,13 @@ def get_provider_info():
170170
"Nebius",
171171
"Ollama",
172172
"OpenAI",
173+
"OpenAI Decisions",
173174
"OpenRouter",
174175
"OVHcloud",
175176
"Pydantic AI Gateway",
176177
"SambaNova",
177178
"Snowflake Cortex",
179+
"System One",
178180
"Together AI",
179181
"TypeSafe",
180182
"Vercel AI Gateway",

‎providers/common/ai/src/airflow/providers/common/ai/utils/sql_validation.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,9 @@ def resolve_sqlglot_dialect(dialect_name: str | None) -> str | None:
6969
exp.Except,
7070
)
7171

72+
# Snowflake ``LIST``/``LS @stage``; sqlglot >= 30.22 only (older releases mis-parse it as exp.Alias).
73+
_LIST_STAGE: type[exp.Expr] | None = getattr(exp, "ListStage", None)
74+
7275
# Read-only metadata statements that introspect the schema without touching data:
7376
# - Describe: DESCRIBE / DESC <table> (and EXPLAIN on some dialects)
7477
# - Show: SHOW TABLES / SHOW COLUMNS / SHOW DATABASES, etc.
@@ -323,10 +326,11 @@ def collect_table_references(
323326
continue
324327
# A bare aliased expression is never a valid top-level statement -- sqlglot only
325328
# produces one by mis-parsing input it does not model. Snowflake ``LIST @stage`` /
326-
# ``LS @stage`` (list a stage's files) parse this way, as ``Column AS Parameter``.
327-
# Read-only mode already rejects it via the statement-type allow-list; this also
328-
# refuses it on the ``allow_writes`` path, where only this scan runs.
329-
if isinstance(stmt, exp.Alias):
329+
# ``LS @stage`` (list a stage's files) parse this way, as ``Column AS Parameter``,
330+
# before sqlglot 30.22, which models them as ``exp.ListStage`` instead.
331+
# Read-only mode already rejects both via the statement-type allow-list; this also
332+
# refuses them on the ``allow_writes`` path, where only this scan runs.
333+
if isinstance(stmt, exp.Alias) or (_LIST_STAGE is not None and isinstance(stmt, _LIST_STAGE)):
330334
unverifiable.append("a statement sqlglot could not parse as a query (e.g. Snowflake LIST @stage)")
331335
continue
332336
# A function whose string argument reaches a file, another table, or a program

‎providers/common/ai/tests/unit/common/ai/operators/test_llm.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -358,7 +358,10 @@ def test_execute_raises_when_cost_cap_exceeded(self, mock_hook_cls, cap_kwargs,
358358
op.execute(context=usage_budget_context)
359359

360360
pushes = {c.kwargs["key"]: c.kwargs["value"] for c in ti.xcom_push.call_args_list}
361-
assert pushes["usage"]["requests"] == 1
361+
# Not `requests`: pydantic-ai >= 2.54 counts only responses the agent acted on, so the
362+
# capped response's request is no longer counted, though its tokens and cost still are.
363+
assert pushes["usage"]["input_tokens"] == 100
364+
assert pushes["usage"]["output_tokens"] == 50
362365
assert pushes["usage"]["cost"] == str(PRICED_COST)
363366

364367
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True)

‎providers/common/ai/tests/unit/common/ai/test_provider_metadata.py‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -81,10 +81,12 @@ def _extract_toolset_block() -> dict:
8181
"nebius": "Nebius",
8282
"ollama": "Ollama",
8383
"openai": "OpenAI",
84+
"openai_decisions": "OpenAI Decisions",
8485
"openrouter": "OpenRouter",
8586
"ovhcloud": "OVHcloud",
8687
"sambanova": "SambaNova",
8788
"snowflake": "Snowflake Cortex",
89+
"system_one": "System One",
8890
"together": "Together AI",
8991
"typesafe": "TypeSafe",
9092
"vercel": "Vercel AI Gateway",
@@ -163,14 +165,15 @@ def _extract_toolset_block() -> dict:
163165
# `provider.yaml`'s vendor list is derived from a current pydantic-ai, but the provider
164166
# supports `pydantic-ai-slim>=2.33.0` (`pyproject.toml`) and CI's "Low dep tests" job installs
165167
# exactly that floor. Vendors upstream added after the floor are legitimately absent there --
166-
# bisected empirically: `vllm` landed in 2.38.0, `github_copilot` in 2.42.0, and `typesafe` in
167-
# 2.45.0 (`openai_codex` also landed after the floor, in 2.41.0, but is excluded from this
168-
# reasoning -- see `OPENAI_CODEX_AVAILABLE_FROM` above) -- so only the "installed but not
169-
# declared" direction holds on every supported version. The opposite direction ("declared but
170-
# not installed") is a real drift signal solely from the release that carries every vendor
171-
# currently declared, i.e. the latest of the three, `typesafe`'s 2.45.0: raise this when a
172-
# newly declared vendor lands in a later release.
173-
ALL_DECLARED_VENDORS_AVAILABLE_FROM = Version("2.45.0")
168+
# bisected empirically: `vllm` landed in 2.38.0, `github_copilot` in 2.42.0, `typesafe` in
169+
# 2.45.0, `system_one` in 2.53.0 and `openai_decisions` in 2.55.0 (`openai_codex` also landed
170+
# after the floor, in 2.41.0, but is excluded from this reasoning -- see
171+
# `OPENAI_CODEX_AVAILABLE_FROM` above) -- so only the "installed but not declared" direction
172+
# holds on every supported version. The opposite direction ("declared but not installed") is a
173+
# real drift signal solely from the release that carries every vendor currently declared, i.e.
174+
# the latest of them, `openai_decisions`'s 2.55.0: raise this when a newly declared vendor lands
175+
# in a later release.
176+
ALL_DECLARED_VENDORS_AVAILABLE_FROM = Version("2.55.0")
174177

175178
requires_all_declared_vendors = pytest.mark.skipif(
176179
Version(Version(version("pydantic-ai-slim")).base_version) < ALL_DECLARED_VENDORS_AVAILABLE_FROM,

‎providers/common/ai/tests/unit/common/ai/utils/test_prompt_cache.py‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,35 @@ def test_setting_names_are_the_defaults_plus_automatic_anthropic_caching(self):
141141
assert {*ALL_DEFAULTS, "anthropic_cache"} == PROMPT_CACHE_SETTING_NAMES
142142

143143

144+
_ANTHROPIC_OK_STREAM = "".join(
145+
f"event: {event['type']}\ndata: {json.dumps(event)}\n\n"
146+
for event in [
147+
{
148+
"type": "message_start",
149+
"message": {
150+
"id": "msg_1",
151+
"type": "message",
152+
"role": "assistant",
153+
"model": "claude-sonnet-4-5",
154+
"content": [],
155+
"stop_reason": None,
156+
"stop_sequence": None,
157+
"usage": {"input_tokens": 10, "output_tokens": 0},
158+
},
159+
},
160+
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
161+
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}},
162+
{"type": "content_block_stop", "index": 0},
163+
{
164+
"type": "message_delta",
165+
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
166+
"usage": {"output_tokens": 1},
167+
},
168+
{"type": "message_stop"},
169+
]
170+
).encode()
171+
172+
144173
class TestPromptCachingOnTheAnthropicWire:
145174
"""What an Anthropic model actually sends, captured at the HTTP transport."""
146175

@@ -149,7 +178,14 @@ def _send(*, model_settings: Any = None, prompt: Any = "hi") -> dict[str, Any]:
149178
bodies: list[dict[str, Any]] = []
150179

151180
def handler(request: httpx2.Request) -> httpx2.Response:
152-
bodies.append(json.loads(request.content))
181+
body = json.loads(request.content)
182+
bodies.append(body)
183+
# pydantic-ai >= 2.52 streams whenever its default max_tokens exceeds the SDK's
184+
# non-streaming limit and the client timeout is not the SDK default, as here.
185+
if body.get("stream"):
186+
return httpx2.Response(
187+
200, headers={"content-type": "text/event-stream"}, content=_ANTHROPIC_OK_STREAM
188+
)
153189
return httpx2.Response(
154190
200,
155191
json={

‎providers/common/ai/tests/unit/common/ai/utils/test_sql_validation.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -418,7 +418,9 @@ def test_legit_cte_reference_is_excluded(self):
418418
def test_flags_inline_comment_as_unverifiable(self, sql, dialect):
419419
"""Comments hide parser-vs-engine differentials (MySQL executable comments), so reject them."""
420420
scan = collect_table_references(parse_sql(sql, dialect=dialect))
421-
assert scan.unverifiable_sources
421+
# Newer sqlglot tokenizes MySQL's `--+1` as arithmetic, as MySQL does (a comment needs
422+
# whitespace after `--`); with no differential left, `secret` is then a scanned table.
423+
assert scan.unverifiable_sources or ("", "", "secret") in scan.tables
422424

423425
@pytest.mark.parametrize(
424426
"sql",

0 commit comments

Comments
 (0)