Fix LLMSQLQueryOperator not stripping single-line markdown code fences (#70137)

* Fix LLMSQLQueryOperator not stripping single-line markdown code fences

_strip_llm_output only removed ```sql fences when the LLM's response
spanned multiple lines. LLMs that wrap short queries in a fence on a
single line (e.g. "```SELECT 1```") passed the fenced text straight
into SQL validation/execution with the backticks still attached. Add
a fallback branch for the single-line case, alongside the existing
multi-line handling.

* Also strip a same-line language tag on single-line fences

Address review feedback: a single-line fence like "```sql SELECT 1```"
left "sql " stuck to the front of the query, since the earlier fix only
stripped the outer backticks. Drop the leading word too, but only when
it's "sql" or the resolved dialect, so a real leading SQL keyword is
never mistaken for a tag.

Author: ColtenOuO

providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py

@@ -152,7 +152,7 @@ def execute(self, context: Context) -> str:
         )
         result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
         log_run_summary(self.log, result)
-        sql = self._strip_llm_output(result.output)
+        sql = self._strip_llm_output(result.output, dialect=self._resolved_dialect)
 
         if self.validate_sql:
             _validate_sql(sql, allowed_types=self.allowed_sql_types, dialect=self._resolved_dialect)
@@ -172,15 +172,25 @@ def execute_complete(self, context: Context, generated_output: str, event: dict[
         return output
 
     @staticmethod
-    def _strip_llm_output(raw: str) -> str:
+    def _strip_llm_output(raw: str, *, dialect: str | None = None) -> str:
         """Strip whitespace and markdown code fences from LLM output."""
         text = raw.strip()
         if text.startswith("```"):
             lines = text.split("\n")
-            # Remove opening fence (```sql, ```, etc.) and closing fence
             if len(lines) >= 2:
+                # Remove opening fence (```sql, ```, etc.) and closing fence
                 end = -1 if lines[-1].strip().startswith("```") else len(lines)
                 text = "\n".join(lines[1:end]).strip()
+            elif text.endswith("```") and len(text) > 6:
+                # Whole fenced block on one line, e.g. "```sql SELECT 1```" -> "SELECT 1".
+                # Only drop the leading word if it's a known tag, so a real keyword
+                # like "SELECT" is never mistaken for one.
+                inner = text[3:-3].strip()
+                tags = {"sql", *([dialect.lower()] if dialect else [])}
+                first_word, sep, rest = inner.partition(" ")
+                if sep and first_word.lower() in tags:
+                    inner = rest.strip()
+                text = inner
         return text
 
     def _get_schema_context(self) -> str:

providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py

@@ -77,11 +77,49 @@ class TestStripLLMOutput:
                 "SELECT 1",
                 id="missing_closing_fence",
             ),
+            pytest.param(
+                "```SELECT 1```",
+                "SELECT 1",
+                id="single_line_fence_no_language_tag",
+            ),
+            pytest.param(
+                "```SELECT * FROM users LIMIT 10```",
+                "SELECT * FROM users LIMIT 10",
+                id="single_line_fence_with_query",
+            ),
+            pytest.param(
+                "```sql SELECT 1```",
+                "SELECT 1",
+                id="single_line_fence_with_language_tag",
+            ),
+            pytest.param(
+                "```SQL SELECT 1```",
+                "SELECT 1",
+                id="single_line_fence_with_uppercase_language_tag",
+            ),
         ),
     )
     def test_strip_llm_output(self, raw, expected):
         assert LLMSQLQueryOperator._strip_llm_output(raw) == expected
 
+    @pytest.mark.parametrize(
+        ("raw", "dialect", "expected"),
+        (
+            pytest.param("```postgres SELECT 1```", "postgres", "SELECT 1", id="dialect_tag_matches"),
+            pytest.param(
+                "```POSTGRES SELECT 1```", "postgres", "SELECT 1", id="dialect_tag_case_insensitive"
+            ),
+            pytest.param(
+                "```mysql SELECT 1```",
+                "postgres",
+                "mysql SELECT 1",
+                id="dialect_tag_mismatch_left_alone",
+            ),
+        ),
+    )
+    def test_strip_llm_output_with_dialect(self, raw, dialect, expected):
+        assert LLMSQLQueryOperator._strip_llm_output(raw, dialect=dialect) == expected
+
 
 class TestLLMSQLQueryOperator:
     def test_inherits_from_llm_operator(self):