refactor: Ensure Celery leverages the Flask-SQLAlchemy session (#26186)

This commit is contained in:
John Bodley
2024-01-17 17:06:22 +13:00
committed by GitHub
parent aaa4a7b371
commit 7af82ae87d
19 changed files with 932 additions and 348 deletions

View File

@@ -23,7 +23,6 @@ from typing import Any, Callable, List, NamedTuple, Optional
from flask_babel import gettext as __
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.orm import Session
with contextlib.suppress(ImportError, RuntimeError): # pyocient may not be installed
# Ensure pyocient inherits Superset's logging level
@@ -372,13 +371,13 @@ class OcientEngineSpec(BaseEngineSpec):
return "DUMMY_VALUE"
@classmethod
def handle_cursor(cls, cursor: Any, query: Query, session: Session) -> None:
def handle_cursor(cls, cursor: Any, query: Query) -> None:
with OcientEngineSpec.query_id_mapping_lock:
OcientEngineSpec.query_id_mapping[query.id] = cursor.query_id
# Add the query id to the cursor
setattr(cursor, "superset_query_id", query.id)
return super().handle_cursor(cursor, query, session)
return super().handle_cursor(cursor, query)
@classmethod
def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool: