fix(20428): Address-Presto/Trino-Poll-Issue-Refactor (#20434)

* fix(20428)-Address-Presto/Trino-Poll-Issue-Refacto r Update linter * Update to only use BaseEngineSpec handle_cursor * Fix CI Co-authored-by: John Bodley <[email protected]> (cherry picked from commit 8b7262f)
apache · Jun 23, 2022 · cb27003 · cb27003
1 parent 7e48de4
commit cb27003
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 5 deletions.
diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py
@@ -949,11 +949,7 @@ def get_create_view(
             sql = f"SHOW CREATE VIEW {schema}.{table}"
             try:
                 cls.execute(cursor, sql)
-                polled = cursor.poll()
 
-                while polled:
-                    time.sleep(0.2)
-                    polled = cursor.poll()
             except DatabaseError:  # not a VIEW
                 return None
             rows = cls.fetch_data(cursor, 1)

diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py
@@ -15,15 +15,18 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
-from typing import Any, Dict, Optional, TYPE_CHECKING
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
 import simplejson as json
 from flask import current_app
+from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.engine.url import URL
+from sqlalchemy.orm import Session
 
 from superset.databases.utils import make_url_safe
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.presto import PrestoEngineSpec
+from superset.models.sql_lab import Query
 from superset.utils import core as utils
 
 if TYPE_CHECKING:
@@ -77,6 +80,37 @@ def modify_url_for_impersonation(
     def get_allow_cost_estimate(cls, extra: Dict[str, Any]) -> bool:
         return True
 
+    @classmethod
+    def get_table_names(
+        cls,
+        database: "Database",
+        inspector: Inspector,
+        schema: Optional[str],
+    ) -> List[str]:
+        return BaseEngineSpec.get_table_names(
+            database=database,
+            inspector=inspector,
+            schema=schema,
+        )
+
+    @classmethod
+    def get_view_names(
+        cls,
+        database: "Database",
+        inspector: Inspector,
+        schema: Optional[str],
+    ) -> List[str]:
+        return BaseEngineSpec.get_view_names(
+            database=database,
+            inspector=inspector,
+            schema=schema,
+        )
+
+    @classmethod
+    def handle_cursor(cls, cursor: Any, query: Query, session: Session) -> None:
+        """Updates progress information"""
+        BaseEngineSpec.handle_cursor(cursor=cursor, query=query, session=session)
+
     @staticmethod
     def get_extra_params(database: "Database") -> Dict[str, Any]:
         """