def test_get_executor_class(self):
with self.mock_all_executors:
test_engine, test2_engine = all_executors.ALL_EXECUTORS
self.assertEqual(
all_executors.get_executor_class("English", "Test"),
test_engine,
)
self.assertEqual(
all_executors.get_executor_class("English", "Test2"),
test2_engine,
)
self.assertEqual(
all_executors.get_executor_class("French", "Test2"),
test2_engine,
)
def _get_executor_and_params_by_engine_id(cls,
engine_id: int,
session=None):
engine = get_query_engine_by_id(engine_id, session=session)
executor_params = engine.get_engine_params()
executor = get_executor_class(engine.language, engine.executor)
return executor, executor_params, engine.to_dict_admin()
def _perform_check(cls, engine_id: int) -> EngineStatus:
with DBSession() as session:
engine = get_query_engine_by_id(engine_id, session=session)
executor_params = engine.get_engine_params()
return check_connection(
get_executor_class(engine.language, engine.executor),
executor_params)
def _get_sqlalchemy_connection(self, session=None):
engine = get_query_engine_by_id(self._engine_id, session=session)
executor = get_executor_class(engine.language, engine.executor)
executor_params = engine.get_engine_params()
client = executor._get_client(executor_params)
if not isinstance(client, SqlAlchemyClient):
raise ValueError(
f"Client instance {client} is not SqlAlchemy Based")
conn = client._engine.connect()
return conn
def test_query_engine_connection(
name,
language,
executor,
executor_params,
feature_params,
):
status_checker = get_engine_checker_class(feature_params["status_checker"])
executor_class = get_executor_class(language, executor)
pseudo_engine_dict = {
"name": name,
"language": language,
"executor": executor,
"executor_params": executor_params,
"feature_params": feature_params,
}
return status_checker.perform_check_with_executor(executor_class,
executor_params,
pseudo_engine_dict)
def __call__(
self, query: str, engine_id: int, uid: int = None, session=None,
):
"""Start the query execution progress. If async then
it just sets up the necessary variables, if sync
then actually execute the query
Args:
query (str): Query getting executed
engine_id (int): The id of the engine
uid (int, optional): User id for proxy user. Defaults to None.
session (SqlAlchemySession, optional): for querying database
Returns:
Any[][]: Returns the result if sync, otherwise None
"""
engine = get_query_engine_by_id(engine_id, session=session)
client_settings = {
**engine.get_engine_params(),
}
if uid:
user = get_user_by_id(uid, session=session)
client_settings["proxy_user"] = user.username
executor = get_executor_class(engine.language, engine.executor)
if executor.SINGLE_QUERY_QUERY_ENGINE():
statements = [query]
else:
statements = get_statements(query)
if len(statements) == 0:
return None # Empty statement, return None
cursor = executor._get_client(client_settings).cursor()
if self._async:
self._async_run(cursor, statements)
return None
else:
return self._sync_run(cursor, statements)
def __call__(
self,
query: str,
engine_id: int,
uid: int = None,
session=None,
):
"""Start the query execution progress. If async then
it just sets up the necessary variables, if sync
then actually execute the query
Args:
query (str): Query getting executed
engine_id (int): The id of the engine
uid (int, optional): User id for proxy user. Defaults to None.
session (SqlAlchemySession, optional): for querying database
Returns:
Any[][]: Returns the result if sync, otherwise None
"""
engine = get_query_engine_by_id(engine_id, session=session)
client_settings = get_client_setting_from_engine(engine,
uid,
session=session)
self.executor = get_executor_class(engine.language, engine.executor)
statements = parse_statement_from_query(self.executor, query)
if len(statements) == 0:
# Empty statement, return None
return None
cursor = self.executor._get_client(client_settings).cursor()
if self._async:
self._async_run(cursor, statements)
return None
else:
return self._sync_run(cursor, statements)
