(superset) branch resilient-cursor-fetch updated (f0c3282faa -> f4a8070432)

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a change to branch resilient-cursor-fetch
in repository https://gitbox.apache.org/repos/asf/superset.git


 discard f0c3282faa fix(drill): no rows returned
     new f4a8070432 fix(drill): no rows returned

This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version.  This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:

 * -- * -- B -- O -- O -- O   (f0c3282faa)
            \
             N -- N -- N   refs/heads/resilient-cursor-fetch (f4a8070432)

You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.

Any revisions marked "omit" are not gone; other references still
refer to them.  Any revisions marked "discard" are gone forever.

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.


Summary of changes:
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(superset) 01/01: fix(drill): no rows returned

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch resilient-cursor-fetch
in repository https://gitbox.apache.org/repos/asf/superset.git

commit f4a8070432f0aa3229021b8a6962f8e13091bb66
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Fri Feb 9 15:54:02 2024 -0500

    fix(drill): no rows returned
---
 setup.py                          |  2 +-
 superset/db_engine_specs/drill.py | 37 ++++++++++++++++++++++++++++++-------
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/setup.py b/setup.py
index e2ee0db7f2..4318ea8860 100644
--- a/setup.py
+++ b/setup.py
@@ -154,7 +154,7 @@ setup(
         ],
         "db2": ["ibm-db-sa>0.3.8, <=0.4.0"],
         "dremio": ["sqlalchemy-dremio>=1.1.5, <1.3"],
-        "drill": ["sqlalchemy-drill==0.1.dev"],
+        "drill": ["sqlalchemy-drill>=1.1.4, <2"],
         "druid": ["pydruid>=0.6.5,<0.7"],
         "duckdb": ["duckdb-engine>=0.9.5, <0.10"],
         "dynamodb": ["pydynamodb>=0.4.2"],
diff --git a/superset/db_engine_specs/drill.py b/superset/db_engine_specs/drill.py
index fb42409b4e..276ff5b185 100644
--- a/superset/db_engine_specs/drill.py
+++ b/superset/db_engine_specs/drill.py
@@ -14,8 +14,11 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+from __future__ import annotations
+
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any
 from urllib import parse
 
 from sqlalchemy import types
@@ -60,8 +63,8 @@ class DrillEngineSpec(BaseEngineSpec):
 
     @classmethod
     def convert_dttm(
-        cls, target_type: str, dttm: datetime, db_extra: Optional[dict[str, Any]] = None
-    ) -> Optional[str]:
+        cls, target_type: str, dttm: datetime, db_extra: dict[str, Any] | None = None
+    ) -> str | None:
         sqla_type = cls.get_sqla_column_type(target_type)
 
         if isinstance(sqla_type, types.Date):
@@ -76,8 +79,8 @@ class DrillEngineSpec(BaseEngineSpec):
         cls,
         uri: URL,
         connect_args: dict[str, Any],
-        catalog: Optional[str] = None,
-        schema: Optional[str] = None,
+        catalog: str | None = None,
+        schema: str | None = None,
     ) -> tuple[URL, dict[str, Any]]:
         if schema:
             uri = uri.set(database=parse.quote(schema.replace(".", "/"), safe=""))
@@ -89,7 +92,7 @@ class DrillEngineSpec(BaseEngineSpec):
         cls,
         sqlalchemy_uri: URL,
         connect_args: dict[str, Any],
-    ) -> Optional[str]:
+    ) -> str | None:
         """
         Return the configured schema.
         """
@@ -97,7 +100,7 @@ class DrillEngineSpec(BaseEngineSpec):
 
     @classmethod
     def get_url_for_impersonation(
-        cls, url: URL, impersonate_user: bool, username: Optional[str]
+        cls, url: URL, impersonate_user: bool, username: str | None
     ) -> URL:
         """
         Return a modified URL with the username set.
@@ -117,3 +120,23 @@ class DrillEngineSpec(BaseEngineSpec):
                 )
 
         return url
+
+    @classmethod
+    def fetch_data(
+        cls,
+        cursor: Any,
+        limit: int | None = None,
+    ) -> list[tuple[Any, ...]]:
+        """
+        Custom `fetch_data` for Drill.
+
+        When no rows are returned, Drill raises a `RuntimeError` with the message
+        "generator raised StopIteration". This method catches the exception and
+        returns an empty list instead.
+        """
+        try:
+            return super().fetch_data(cursor, limit)
+        except RuntimeError as ex:
+            if str(ex) == "generator raised StopIteration":
+                return []
+            raise