potiuk commented on code in PR #66658:
URL: https://github.com/apache/airflow/pull/66658#discussion_r3215590420
##########
providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py:
##########
@@ -99,7 +106,29 @@ def fetchone(self):
return None
Review Comment:
Now that `fetchmany` consumes, `fetchone` (just above) looks more obviously
broken in comparison — `return self.rows[0]` returns row 0 every time without
advancing. Pre-existing bug, not introduced here, but mixing `fetchone()` and
`fetchmany()` calls will produce surprising results for any caller. Worth
fixing in this PR (it's ~4 lines):
```python
def fetchone(self):
if self._rows:
return self._rows.pop(0)
if self.cursor:
self.execute(statement=self.body["query"])
return self.fetchone()
return None
```
Or file a follow-up issue and link it from the PR body so it doesn't get
forgotten.
##########
providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py:
##########
@@ -99,7 +106,29 @@ def fetchone(self):
return None
def fetchmany(self, size: int | None = None):
- raise NotImplementedError()
+ size = size or self.body["fetch_size"]
+
+ results: list[list[Any]] = []
+
+ while len(results) < size:
+ # Consume rows from the current page buffer before
+ # requesting additional cursor pages from Elasticsearch.
+ if self._rows:
+ remaining = size - len(results)
+
+ results.extend(self._rows[:remaining])
+
+ self._rows = self._rows[remaining:]
Review Comment:
Slicing reassigns and copies the tail every iteration — O(n) per drain. For
typical `fetch_size`s (1000+) and reasonable `fetchmany(size)` it's fine, but a
caller doing `fetchmany(50000)` with `fetch_size=10` gets O(n²) behaviour.
If you want to make this defensive: switch `self._rows` to
`collections.deque` and use `popleft()` — O(1) amortised. Not a blocker for
normal use.
##########
providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py:
##########
@@ -99,7 +106,29 @@ def fetchone(self):
return None
def fetchmany(self, size: int | None = None):
- raise NotImplementedError()
+ size = size or self.body["fetch_size"]
+
+ results: list[list[Any]] = []
+
+ while len(results) < size:
+ # Consume rows from the current page buffer before
+ # requesting additional cursor pages from Elasticsearch.
+ if self._rows:
+ remaining = size - len(results)
+
+ results.extend(self._rows[:remaining])
+
+ self._rows = self._rows[remaining:]
+
+ if len(results) >= size:
+ break
+
+ if not self.cursor:
+ break
+
+ self.execute(statement=self.body["query"])
Review Comment:
Re-entering `execute` here passes the original query (and whatever `params`
were stored in `self.body`) alongside the cursor. The Elasticsearch SQL API
expects **only** `cursor` on pagination requests; sending `query` and `params`
is at best ignored and at worst rejected depending on ES version.
Pre-existing in `fetchall` (same pattern), so not introduced by this PR —
but now `fetchmany` exercises this path too. Worth a separate fix to `execute`
that sends only `{"cursor": ...}` when the body has a cursor. Could be
follow-up.
##########
providers/elasticsearch/src/airflow/providers/elasticsearch/hooks/elasticsearch.py:
##########
@@ -70,7 +74,7 @@ def cursor(self):
@property
def rows(self):
- return self.response.get("rows", [])
+ return self._rows
Review Comment:
The `rows` property now returns the mutable buffer instead of reading from
`self.response`, which changes `rowcount`'s semantics: it used to mean "rows in
the most recent response page" and now means "rows remaining in the buffer".
Callers who relied on `rowcount` for total result-set size will see decreasing
values as `fetchmany` drains.
DB-API 2.0 says `rowcount` is "the number of rows the last `.execute*()`
produced", which doesn't cleanly fit either interpretation. Consider either:
- Keeping `rowcount` reading from `self.response.get("rows", [])` (preserves
old semantics for that one property), or
- Documenting the change in the PR body so downstream users know.
The test was updated from `len(ROWS)` to `len(ROWS_PAGE_1)` — which papers
over the change rather than flagging it.
##########
providers/elasticsearch/tests/unit/elasticsearch/hooks/test_elasticsearch.py:
##########
@@ -109,22 +126,72 @@ def test_fetchone(self):
assert cursor.fetchone() == ROWS[0]
- def test_fetchmany(self):
+ @pytest.mark.parametrize(
+ ("size", "expected"),
+ [
+ (1, ROWS[:1]),
+ (2, ROWS[:2]),
+ (5, ROWS[:5]),
+ ],
+ )
+ def test_fetchmany(self, size, expected):
cursor = ElasticsearchSQLCursor(es=self.es, options={})
cursor.execute("SELECT * FROM hollywood.actors")
- with pytest.raises(NotImplementedError):
- cursor.fetchmany()
+ records = cursor.fetchmany(size)
+
+ assert records == expected
+
+ def test_fetchmany_consumes_rows(self):
+ cursor = ElasticsearchSQLCursor(es=self.es, options={})
+ cursor.execute("SELECT * FROM hollywood.actors")
+
+ first_batch = cursor.fetchmany(2)
+ second_batch = cursor.fetchmany(2)
+
+ assert first_batch == ROWS[:2]
+ assert second_batch == ROWS[2:4]
+
+ def test_fetchmany_exhausts_rows(self):
+ cursor = ElasticsearchSQLCursor(es=self.es, options={})
+ cursor.execute("SELECT * FROM hollywood.actors")
+
+ records = cursor.fetchmany(100)
+
+ assert records == ROWS
+
+ # Further calls should return empty list.
+ assert cursor.fetchmany(10) == []
+
+ def test_fetchmany_uses_default_fetch_size(self):
+ cursor = ElasticsearchSQLCursor(es=self.es, fetch_size=2)
+ cursor.execute("SELECT * FROM hollywood.actors")
+
+ records = cursor.fetchmany()
+
+ assert records == ROWS[:2]
def test_fetchall(self):
Review Comment:
Missing case: single-page result with no cursor returned. All current tests
use the paginated response setup (`RESPONSES = [RESPONSE,
RESPONSE_WITHOUT_CURSOR]`), so the path where `execute` returns rows without a
cursor and `fetchmany(size > rowcount)` short-circuits on the `if not
self.cursor: break` branch isn't exercised.
```python
def test_fetchmany_single_page(self, ...):
# mock es.sql.query to return RESPONSE_WITHOUT_CURSOR only
cursor = ElasticsearchSQLCursor(es=self.es, options={})
cursor.execute("SELECT ...")
assert cursor.fetchmany(100) == ROWS_PAGE_2
assert cursor.fetchmany(10) == []
```
--
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]
