This is an automated email from the ASF dual-hosted git repository.
michaelsmolina pushed a commit to branch 4.1-airbnb
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/4.1-airbnb by this push:
new 4d2f13ef1a fix(trino): db session error in handle cursor (#31265)
4d2f13ef1a is described below
commit 4d2f13ef1a8c6d8b8647c9c5818cc879d4264940
Author: JUST.in DO IT <[email protected]>
AuthorDate: Tue Dec 3 11:57:37 2024 -0800
fix(trino): db session error in handle cursor (#31265)
(cherry picked from commit 1e0c04fc15477aaa2ca144302e6c09f5ba0e39ff)
---
superset/db_engine_specs/trino.py | 10 +++++++---
1 file changed, 7 insertions(+), 3 deletions(-)
diff --git a/superset/db_engine_specs/trino.py
b/superset/db_engine_specs/trino.py
index e80e4de47c..3cbf6ebe79 100644
--- a/superset/db_engine_specs/trino.py
+++ b/superset/db_engine_specs/trino.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=consider-using-transaction
from __future__ import annotations
import contextlib
@@ -185,6 +184,8 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
if tracking_url := cls.get_tracking_url(cursor):
query.tracking_url = tracking_url
+ db.session.commit() # pylint: disable=consider-using-transaction
+
# if query cancelation was requested prior to the handle_cursor call,
but
# the query was still executed, trigger the actual query cancelation
now
if query.extra.get(QUERY_EARLY_CANCEL_KEY):
@@ -213,6 +214,7 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
# Fetch the query ID beforehand, since it might fail inside the thread
due to
# how the SQLAlchemy session is handled.
query_id = query.id
+ query_database = query.database
execute_result: dict[str, Any] = {}
execute_event = threading.Event()
@@ -234,7 +236,7 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
with app.app_context():
for key, value in g_copy.__dict__.items():
setattr(g, key, value)
- cls.execute(cursor, sql, query.database)
+ cls.execute(cursor, sql, query_database)
except Exception as ex: # pylint: disable=broad-except
results["error"] = ex
finally:
@@ -251,6 +253,8 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
)
execute_thread.start()
+ # Wait for the thread to start before continuing
+ time.sleep(0.1)
# Wait for a query ID to be available before handling the cursor, as
# it's required by that method; it may never become available on error.
while not cursor.query_id and not execute_event.is_set():
@@ -272,7 +276,7 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
def prepare_cancel_query(cls, query: Query) -> None:
if QUERY_CANCEL_KEY not in query.extra:
query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
- db.session.commit()
+ db.session.commit() # pylint: disable=consider-using-transaction
@classmethod
def cancel_query(cls, cursor: Cursor, query: Query, cancel_query_id: str)
-> bool:
