kaxil commented on code in PR #74017:
URL: https://github.com/apache/airflow/pull/74017#discussion_r4157357317
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
+ ),
+ }
+ )
+
+ full: dict[str, Any] = {"columns": selected, "column_count": len(selected)}
+ if name_contains is not None:
+ full["name_contains"] = name_contains
+ full["total_columns"] = total_columns
+ if len(selected) <= max_columns and _size(full) <= max_result_bytes:
+ return _dumps(full)
+
+ return _summarize_schema(
+ selected,
+ total_columns=total_columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+
+
+def _summarize_schema(
+ selected: list[dict[str, str]],
+ *,
+ total_columns: int,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None,
+) -> str:
+ """Build the bounded summary returned when the full column list does not
fit."""
+ # Count is checked before bytes: it is the cheaper, more explainable bound
and the one the
+ # issue is about. A bytes-only truncation then means "count fits but
names/types are
+ # pathologically long", a distinct and rarer signal.
+ truncated_by = "max_columns" if len(selected) > max_columns else
"max_result_bytes"
+ hint = (
+ f"This filter matches {len(selected)} columns, more than can be
returned at once. "
+ "Use a more specific name_contains substring to narrow to the columns
you need."
+ if name_contains is not None
+ else (
+ f"This table has {len(selected)} columns, more than can be
returned at once. Call "
+ "get_schema again with a name_contains substring to return only
the matching columns."
+ )
+ )
+
+ output: dict[str, Any] = {
+ "column_count": len(selected),
+ "truncated": True,
+ "truncated_by": truncated_by,
+ "hint": hint,
+ }
+ if name_contains is not None:
+ output["name_contains"] = name_contains
+ output["total_columns"] = total_columns
+
+ # The core above is the guaranteed-useful payload. Add the histogram and a
sample of columns
+ # only while they fit, reserving the core first -- the same
contiguous-prefix accounting
+ # build_query_result uses for rows -- so a pathologically small budget
still returns the core.
+ output["type_histogram"] = _build_type_histogram(selected)
+ output["sample_columns"] = []
+ if _size(output) > max_result_bytes:
+ del output["type_histogram"]
+ return _dumps(output)
+ budget = max_result_bytes - _size(output)
+ for col in selected[:max_columns]:
Review Comment:
With the default `max_columns=100` the sample is the first 100 columns, so
the summary is about as big as the list it replaces. In a quick probe with
`customer_attribute_N` names, a 101-column table returned 5270 bytes of summary
where the full list was 5073, and column 101 was gone. A 1000-column table
whose full list is ~51 KB (under the 64 KiB budget) comes back as the first 100
columns plus a hint, which is the "truncate to the first N" behaviour the new
docs section says this design avoids.
Would a small fixed preview (say 20 columns) work better here, so the
summary is actually smaller than what it replaces? Or, if 100 isn't meant as a
context policy of its own, drop the count trigger and let `max_result_bytes`
decide. `test_too_many_columns_are_summarized_not_listed` pins
`len(sample_columns) == 100`, so it would move with whichever you pick.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
+ ),
+ }
+ )
+
+ full: dict[str, Any] = {"columns": selected, "column_count": len(selected)}
+ if name_contains is not None:
+ full["name_contains"] = name_contains
+ full["total_columns"] = total_columns
+ if len(selected) <= max_columns and _size(full) <= max_result_bytes:
+ return _dumps(full)
+
+ return _summarize_schema(
+ selected,
+ total_columns=total_columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+
+
+def _summarize_schema(
+ selected: list[dict[str, str]],
+ *,
+ total_columns: int,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None,
+) -> str:
+ """Build the bounded summary returned when the full column list does not
fit."""
+ # Count is checked before bytes: it is the cheaper, more explainable bound
and the one the
+ # issue is about. A bytes-only truncation then means "count fits but
names/types are
+ # pathologically long", a distinct and rarer signal.
+ truncated_by = "max_columns" if len(selected) > max_columns else
"max_result_bytes"
+ hint = (
+ f"This filter matches {len(selected)} columns, more than can be
returned at once. "
+ "Use a more specific name_contains substring to narrow to the columns
you need."
+ if name_contains is not None
+ else (
+ f"This table has {len(selected)} columns, more than can be
returned at once. Call "
+ "get_schema again with a name_contains substring to return only
the matching columns."
+ )
+ )
+
+ output: dict[str, Any] = {
+ "column_count": len(selected),
+ "truncated": True,
+ "truncated_by": truncated_by,
+ "hint": hint,
+ }
+ if name_contains is not None:
+ output["name_contains"] = name_contains
+ output["total_columns"] = total_columns
+
+ # The core above is the guaranteed-useful payload. Add the histogram and a
sample of columns
+ # only while they fit, reserving the core first -- the same
contiguous-prefix accounting
+ # build_query_result uses for rows -- so a pathologically small budget
still returns the core.
+ output["type_histogram"] = _build_type_histogram(selected)
+ output["sample_columns"] = []
+ if _size(output) > max_result_bytes:
+ del output["type_histogram"]
+ return _dumps(output)
Review Comment:
If the histogram alone doesn't fit, this returns before the sample loop
runs, so every other column loses its preview too. Top-K caps the number of
keys but not their length: 150 VARCHAR columns plus one struct column with a
few thousand nested fields comes back with `sample_columns: []` at the default
budget, and a lowered `max_result_bytes` hits it with much smaller types.
Deleting the histogram and falling through to the loop would still fill the
sample from what's left.
##########
providers/common/ai/docs/toolsets/sql.rst:
##########
@@ -205,6 +206,8 @@ Parameters
transferred is its own call. See :ref:`bounded-query-results`.
- ``max_result_bytes``: Budget for the serialized ``query`` result. Default 64
KiB.
Review Comment:
`max_result_bytes` now also triggers the `get_schema` summary, but this
still calls it the budget for the `query` result. Same wording in
`datafusion.rst` and in the `max_result_bytes` docstrings of `SQLToolset` and
`DataFusionToolset`.
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py:
##########
@@ -187,16 +188,48 @@ def test_returns_column_info(self):
result = asyncio.run(
ts.call_tool("get_schema", {"table_name": "users"},
ctx=MagicMock(), tool=MagicMock())
)
- columns = json.loads(result)
- assert columns == [{"name": "id", "type": "INTEGER"}, {"name": "name",
"type": "VARCHAR"}]
+ data = json.loads(result)
+ assert data == {
+ "columns": [{"name": "id", "type": "INTEGER"}, {"name": "name",
"type": "VARCHAR"}],
+ "column_count": 2,
+ }
mock_hook.get_table_schema.assert_called_once_with("users",
schema=None)
+ def test_name_contains_filters_the_columns(self):
+ """``name_contains`` threads from the tool call through to the bounded
result."""
+ ts = SQLToolset("pg_default")
+ ts._hook = _make_mock_db_hook(
+ table_schema=[
+ {"name": "id", "type": "INTEGER"},
+ {"name": "customer_name", "type": "VARCHAR"},
+ ]
+ )
+
+ result = asyncio.run(
+ ts.call_tool(
+ "get_schema",
+ {"table_name": "users", "name_contains": "name"},
+ ctx=MagicMock(),
+ tool=MagicMock(),
+ )
+ )
+ data = json.loads(result)
+ assert data["columns"] == [{"name": "customer_name", "type":
"VARCHAR"}]
+ assert data["name_contains"] == "name"
+ assert data["total_columns"] == 2
+
def test_blocks_table_not_in_allowed_list(self):
+ """The allow-list guard fires before any introspection or filtering."""
Review Comment:
The docstring says the guard fires before any introspection, but nothing
asserts that. `ts._hook.get_table_schema.assert_not_called()` would pin it.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
Review Comment:
On a table wider than `max_columns`, following this hint gets the agent a
summary whose hint sends it back to `name_contains`, so the two point at each
other (tried with 3000 columns and `name_contains="zzz"`). Could this suggest a
shorter or different substring when `total_columns > max_columns`? This branch
also leaves out `total_columns`, which `sql.rst` says every filtered call adds.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
+ ),
+ }
+ )
+
+ full: dict[str, Any] = {"columns": selected, "column_count": len(selected)}
+ if name_contains is not None:
+ full["name_contains"] = name_contains
+ full["total_columns"] = total_columns
+ if len(selected) <= max_columns and _size(full) <= max_result_bytes:
+ return _dumps(full)
+
+ return _summarize_schema(
+ selected,
+ total_columns=total_columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+
+
+def _summarize_schema(
+ selected: list[dict[str, str]],
+ *,
+ total_columns: int,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None,
+) -> str:
+ """Build the bounded summary returned when the full column list does not
fit."""
+ # Count is checked before bytes: it is the cheaper, more explainable bound
and the one the
+ # issue is about. A bytes-only truncation then means "count fits but
names/types are
Review Comment:
nit: there's no issue linked to this PR (#71317 is about `query`), so "the
one the issue is about" won't mean anything to a later reader.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
+ ),
+ }
+ )
+
+ full: dict[str, Any] = {"columns": selected, "column_count": len(selected)}
+ if name_contains is not None:
+ full["name_contains"] = name_contains
+ full["total_columns"] = total_columns
+ if len(selected) <= max_columns and _size(full) <= max_result_bytes:
+ return _dumps(full)
+
+ return _summarize_schema(
+ selected,
+ total_columns=total_columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+
+
+def _summarize_schema(
+ selected: list[dict[str, str]],
+ *,
+ total_columns: int,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None,
+) -> str:
+ """Build the bounded summary returned when the full column list does not
fit."""
+ # Count is checked before bytes: it is the cheaper, more explainable bound
and the one the
+ # issue is about. A bytes-only truncation then means "count fits but
names/types are
+ # pathologically long", a distinct and rarer signal.
+ truncated_by = "max_columns" if len(selected) > max_columns else
"max_result_bytes"
+ hint = (
+ f"This filter matches {len(selected)} columns, more than can be
returned at once. "
+ "Use a more specific name_contains substring to narrow to the columns
you need."
+ if name_contains is not None
+ else (
+ f"This table has {len(selected)} columns, more than can be
returned at once. Call "
Review Comment:
When `truncated_by` is `max_result_bytes` the count fits, so "more than can
be returned at once" isn't the reason, and one 70 KB column name gives "This
table has 1 columns". Maybe key the wording on `truncated_by` and reuse the
`plural` logic from the no-match branch?
##########
providers/common/ai/docs/changelog.rst:
##########
@@ -39,6 +39,16 @@ Changelog
attempt is checked and counted on its own, unchanged. See :ref:`the
cross-attempt usage
budget <agent-usage-budget>`.
+.. note::
+ ``get_schema`` on ``SQLToolset`` and ``DataFusionToolset`` now returns a
JSON object
+ ``{"columns": [{"name", "type"}, ...], "column_count": N}`` instead of a
bare JSON array of
+ columns. The tool also accepts an optional ``name_contains`` substring
filter, and on a table
+ with more columns than ``max_columns`` (default 100), or one whose
serialized columns exceed
+ ``max_result_bytes``, it returns a bounded summary (``column_count``, a
``type_histogram`` and a
+ ``sample_columns`` preview, with ``truncated``, ``truncated_by`` and a
``hint``) in place of the
+ full list. Update any system prompt or direct ``call_tool("get_schema",
...)`` caller that read
+ the old top-level array to read ``result["columns"]`` instead.
Review Comment:
`result["columns"]` raises `KeyError` on a summary, which by default is any
table over 100 columns. Worth telling callers to check `truncated` first, since
a summary carries `sample_columns` and no `columns` key?
##########
providers/common/ai/tests/unit/common/ai/utils/test_query_results.py:
##########
@@ -169,3 +169,99 @@ def
test_a_secret_that_json_escapes_is_masked_in_the_rows(register_secret):
result = _build(["user", "password"], [["admin", secret]])
assert result["rows"] == [["admin", "***"]]
+
+
+def _schema(columns, *, max_columns=100, max_result_bytes=65_536,
name_contains=None) -> dict:
+ return json.loads(
+ build_schema_result(
+ columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+ )
+
+
+def _cols(n: int, *, type_: str = "VARCHAR", prefix: str = "col") ->
list[dict[str, str]]:
+ return [{"name": f"{prefix}_{i}", "type": type_} for i in range(n)]
+
+
+class TestSchemaResult:
+ def test_small_table_returns_every_column(self):
+ cols = _cols(3)
+ assert _schema(cols) == {"columns": cols, "column_count": 3}
+
+ def test_name_contains_filters_case_insensitively(self):
+ cols = [
+ {"name": "CustomerId", "type": "INT"},
+ {"name": "amount", "type": "NUMERIC"},
+ {"name": "customer_name", "type": "VARCHAR"},
+ ]
+ data = _schema(cols, name_contains="customer")
+ assert data["columns"] == [cols[0], cols[2]]
+ assert data["column_count"] == 2
+ assert data["name_contains"] == "customer"
+ assert data["total_columns"] == 3
+ assert "truncated" not in data
+
+ def test_empty_name_contains_is_treated_as_no_filter(self):
+ cols = _cols(3)
+ assert _schema(cols, name_contains="") == {"columns": cols,
"column_count": 3}
+
+ def test_name_contains_with_no_matches_guides_without_erroring(self):
+ data = _schema(_cols(5), name_contains="zzz")
+ assert data["columns"] == []
+ assert data["column_count"] == 0
+ assert data["name_contains"] == "zzz"
+ assert "error" not in data
+ assert "all 5 columns" in data["hint"]
+
+ def test_too_many_columns_are_summarized_not_listed(self):
+ data = _schema(_cols(250), max_columns=100)
+ assert data["truncated"] is True
+ assert data["truncated_by"] == "max_columns"
+ assert data["column_count"] == 250
+ assert "columns" not in data
+ assert data["type_histogram"] == {"VARCHAR": 250}
+ assert len(data["sample_columns"]) == 100
+ assert data["sample_columns"][0] == {"name": "col_0", "type":
"VARCHAR"}
+ assert "name_contains" in data["hint"]
+
+ def test_long_names_blow_the_byte_budget_despite_few_columns(self):
+ cols = [{"name": "x" * 500, "type": "VARCHAR"} for _ in range(10)]
+ data = _schema(cols, max_columns=100, max_result_bytes=512)
+ assert data["truncated"] is True
+ assert data["truncated_by"] == "max_result_bytes"
Review Comment:
None of the summary tests check the serialized result against the budget
(`len(raw.encode()) <= max_result_bytes`), which is the contract the docstring
promises. Nothing checks that `self._max_columns` / `self._max_result_bytes`
reach `build_schema_result` from either toolset either: passing the module
defaults there instead still passes every test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
