This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 005eadcacd8 fix(opensearch): page drill-to-detail samples with the
OpenSearch SQL response format (#44703)
005eadcacd8 is described below
commit 005eadcacd8712946c7b8f940ac3e73b50455a4a
Author: Amin Ghadersohi <[email protected]>
AuthorDate: Wed Sep 30 15:17:05 2026 +1000
fix(opensearch): page drill-to-detail samples with the OpenSearch SQL
response format (#44703)
---
superset/db_engine_specs/elasticsearch.py | 38 ++++--
.../db_engine_specs/test_elasticsearch.py | 144 ++++++++++++++++++++-
2 files changed, 172 insertions(+), 10 deletions(-)
diff --git a/superset/db_engine_specs/elasticsearch.py
b/superset/db_engine_specs/elasticsearch.py
index c295071951d..048e086b7dd 100644
--- a/superset/db_engine_specs/elasticsearch.py
+++ b/superset/db_engine_specs/elasticsearch.py
@@ -45,6 +45,10 @@ def _fetch_page_via_cursor(
page_size: int,
sql_path: str,
close_path: str,
+ *,
+ headers: dict[str, str] | None = None,
+ rows_key: str = "rows",
+ columns_key: str = "columns",
) -> tuple[list[list[Any]], list[str]]:
"""
Iterate Elasticsearch/OpenSearch SQL cursor pagination to return a single
@@ -63,6 +67,10 @@ def _fetch_page_via_cursor(
Deep pagination (hundreds of pages) will therefore be noticeably slower
than on ``OFFSET``-capable engines. This is a protocol limitation, not
an implementation choice.
+
+ ``headers`` are sent with every request when given. ``rows_key`` and
+ ``columns_key`` name the response fields: Elasticsearch answers
+ ``rows``/``columns``, the OpenSearch SQL plugin ``datarows``/``schema``.
"""
# The Elasticsearch SQL API rejects trailing semicolons, and any LIMIT
# in the submitted statement caps the result set before the cursor can
@@ -75,19 +83,23 @@ def _fetch_page_via_cursor(
r"\s+LIMIT\s+\d+\s*$", "", sanitized_sql, flags=re.IGNORECASE
)
- # The raw transport does not auto-set Content-Type the way the Python
- # DB-API driver does; ES rejects POSTs without a JSON content type.
- json_headers = {"Content-Type": "application/json"}
+ request_kwargs: dict[str, Any] = {} if headers is None else {"headers":
headers}
with database.get_raw_connection() as conn:
transport = conn.es.transport
response = transport.perform_request(
"POST",
sql_path,
- headers=json_headers,
body={"query": sanitized_sql, "fetch_size": page_size},
+ **request_kwargs,
)
- columns = [col["name"] for col in response.get("columns", [])]
- rows = response.get("rows", [])
+ # Column metadata comes from the remote service; fall back to a
+ # positional label rather than failing on an entry without a name.
+ # OpenSearch schema supplies aliases; Elasticsearch supplies names
only.
+ columns = [
+ col.get("alias") or col.get("name") or f"column_{idx}"
+ for idx, col in enumerate(response.get(columns_key, []))
+ ]
+ rows = response.get(rows_key, [])
cursor = response.get("cursor")
try:
@@ -100,10 +112,10 @@ def _fetch_page_via_cursor(
response = transport.perform_request(
"POST",
sql_path,
- headers=json_headers,
body={"cursor": cursor},
+ **request_kwargs,
)
- rows = response.get("rows", [])
+ rows = response.get(rows_key, [])
cursor = response.get("cursor")
return rows, columns
@@ -115,8 +127,8 @@ def _fetch_page_via_cursor(
transport.perform_request(
"POST",
close_path,
- headers=json_headers,
body={"cursor": cursor},
+ **request_kwargs,
)
except Exception: # pylint: disable=broad-except
logger.warning(
@@ -253,6 +265,9 @@ class ElasticSearchEngineSpec(BaseEngineSpec): # pylint:
disable=abstract-metho
page_size=page_size,
sql_path=cls.SQL_ENDPOINT,
close_path=cls.SQL_CLOSE_ENDPOINT,
+ # elasticsearch-py's raw transport does not set Content-Type the
+ # way the DB-API driver does; ES rejects POSTs without it.
+ headers={"Content-Type": "application/json"},
)
@classmethod
@@ -367,6 +382,11 @@ class OpenDistroEngineSpec(BaseEngineSpec): # pylint:
disable=abstract-method
page_size=page_size,
sql_path=cls.SQL_ENDPOINT,
close_path=cls.SQL_CLOSE_ENDPOINT,
+ # opensearch-py already sends Content-Type: adding it again makes
+ # OpenSearch reject the request ("only one Content-Type header
+ # should be provided"). The SQL plugin answers in its JDBC format.
+ rows_key="datarows",
+ columns_key="schema",
)
@classmethod
diff --git a/tests/unit_tests/db_engine_specs/test_elasticsearch.py
b/tests/unit_tests/db_engine_specs/test_elasticsearch.py
index ef0fa4070c1..49a8353d9ef 100644
--- a/tests/unit_tests/db_engine_specs/test_elasticsearch.py
+++ b/tests/unit_tests/db_engine_specs/test_elasticsearch.py
@@ -450,7 +450,7 @@ def
test_opendistro_fetch_data_with_cursor_uses_opendistro_endpoints() -> None:
database = _build_fake_database(
[
- {"columns": [{"name": "a"}], "rows": [[42]], "cursor": "OD-1"},
+ {"schema": [{"name": "a"}], "datarows": [[42]], "cursor": "OD-1"},
{}, # close
]
)
@@ -468,3 +468,145 @@ def
test_opendistro_fetch_data_with_cursor_uses_opendistro_endpoints() -> None:
calls = database._transport.perform_request.call_args_list
assert calls[0][0][1] == "/_opendistro/_sql"
assert calls[1][0][1] == "/_opendistro/_sql/close"
+
+
+def test_opendistro_fetch_data_with_cursor_reads_jdbc_format_pages() -> None:
+ """
+ The OpenSearch SQL plugin answers in its JDBC format: ``schema`` and
+ ``datarows`` on the first page, ``datarows`` on follow-up pages. Reading
+ Elasticsearch's ``columns``/``rows`` keys returned an empty page.
+ """
+ from superset.db_engine_specs.elasticsearch import OpenDistroEngineSpec
+
+ database = _build_fake_database(
+ [
+ {
+ "schema": [{"name": "a"}, {"name": "COUNT(*)", "alias": "c"}],
+ "datarows": [[0, 1]],
+ "cursor": "OD-1",
+ },
+ {"datarows": [[1, 2]], "cursor": "OD-2"},
+ {}, # close
+ ]
+ )
+
+ rows, cols = OpenDistroEngineSpec.fetch_data_with_cursor(
+ database=database,
+ sql="SELECT a, COUNT(*) AS c FROM idx GROUP BY a",
+ page_index=1,
+ page_size=1,
+ )
+
+ assert cols == ["a", "c"]
+ assert rows == [[1, 2]]
+
+
+def test_opendistro_fetch_data_with_cursor_sends_no_extra_content_type() ->
None:
+ """
+ opensearch-py sets Content-Type itself; a second one makes OpenSearch
+ reject the request with "only one Content-Type header should be provided".
+ """
+ from superset.db_engine_specs.elasticsearch import OpenDistroEngineSpec
+
+ database = _build_fake_database(
+ [
+ {"schema": [{"name": "a"}], "datarows": [[0]], "cursor": "OD-1"},
+ {"datarows": [[1]], "cursor": "OD-2"},
+ {}, # close
+ ]
+ )
+
+ OpenDistroEngineSpec.fetch_data_with_cursor(
+ database=database,
+ sql="SELECT a FROM idx",
+ page_index=1,
+ page_size=1,
+ )
+
+ calls = database._transport.perform_request.call_args_list
+ assert len(calls) == 3
+ for call in calls:
+ assert "headers" not in call.kwargs
+
+
+def test_fetch_data_with_cursor_tolerates_column_without_name() -> None:
+ """
+ Column metadata comes from the remote service: an entry without ``name``
+ must not raise KeyError. Alias wins, then name, then a positional label.
+ """
+ from superset.db_engine_specs.elasticsearch import ElasticSearchEngineSpec
+
+ database = _build_fake_database(
+ [
+ {
+ "columns": [{"name": "a"}, {"alias": "b"}, {"type": "long"}],
+ "rows": [[1, 2, 3]],
+ }
+ ]
+ )
+
+ rows, cols = ElasticSearchEngineSpec.fetch_data_with_cursor(
+ database=database,
+ sql="SELECT a, b, c FROM idx",
+ page_index=0,
+ page_size=10,
+ )
+
+ assert cols == ["a", "b", "column_2"]
+ assert rows == [[1, 2, 3]]
+
+
+def test_fetch_page_via_cursor_forwards_explicit_empty_headers() -> None:
+ """
+ An explicit empty ``headers`` dict is forwarded as given; only ``None``
+ omits the kwarg.
+ """
+ from superset.db_engine_specs.elasticsearch import _fetch_page_via_cursor
+
+ database = _build_fake_database([{"columns": [{"name": "a"}], "rows":
[[0]]}])
+
+ _fetch_page_via_cursor(
+ database=database,
+ sql="SELECT a FROM idx",
+ page_index=0,
+ page_size=10,
+ sql_path="/_sql",
+ close_path="/_sql/close",
+ headers={},
+ )
+
+ calls = database._transport.perform_request.call_args_list
+ assert calls[0].kwargs["headers"] == {}
+
+
[email protected]("page_index, expected_rows", [(2, [[2]]), (3, [])])
+def test_opendistro_cursor_later_pages_and_exhaustion(
+ page_index: int, expected_rows: list[list[int]]
+) -> None:
+ """Preserve JDBC column aliases across later pages and cursor
exhaustion."""
+ from superset.db_engine_specs.elasticsearch import OpenDistroEngineSpec
+
+ database = _build_fake_database(
+ [
+ {
+ "schema": [{"name": "COUNT(*)", "alias": "c"}],
+ "datarows": [[0]],
+ "cursor": "OD-1",
+ },
+ {"datarows": [[1]], "cursor": "OD-2"},
+ {"datarows": [[2]]},
+ ]
+ )
+ rows, columns = OpenDistroEngineSpec.fetch_data_with_cursor(
+ database=database,
+ sql="SELECT COUNT(*) AS c FROM idx",
+ page_index=page_index,
+ page_size=1,
+ )
+ assert rows == expected_rows
+ assert columns == ["c"]
+ calls = database._transport.perform_request.call_args_list
+ assert len(calls) == 3
+ assert calls[1].kwargs["body"] == {"cursor": "OD-1"}
+ assert calls[2].kwargs["body"] == {"cursor": "OD-2"}
+ assert all("headers" not in call.kwargs for call in calls)
