Skip to content

Commit

Permalink
fix(20428): Address-Presto/Trino-Poll-Issue-Refactor (#20434)
Browse files Browse the repository at this point in the history
* 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 <4567245+john-bodley@users.noreply.github.com>
(cherry picked from commit 8b7262f)
  • Loading branch information
Thelin90 authored and michael-s-molina committed Jun 23, 2022
1 parent 7e48de4 commit cb27003
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 5 deletions.
4 changes: 0 additions & 4 deletions superset/db_engine_specs/presto.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
36 changes: 35 additions & 1 deletion superset/db_engine_specs/trino.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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]:
"""
Expand Down

0 comments on commit cb27003

Please sign in to comment.