def create_query_snippet(
context=None,
title=None,
engine_id=None,
description=None,
is_public=None,
golden=None,
):
created_by = current_user.id
api_assert(len(context) > 0, "No empty context")
api_assert(len(title) > 0, "No empty title")
api_assert(engine_id is not None, "Must specify engine")
api_assert(
not golden or current_user.is_admin, "Only data gurus can create golden snippet"
)
with DBSession() as session:
verify_query_engine_permission(engine_id, session=session)
return logic.create_snippet(
created_by,
context=context,
title=title,
engine_id=engine_id,
description=description,
is_public=is_public,
golden=golden,
session=session,
)
def create_query_execution(query,
engine_id,
data_cell_id=None,
originator=None):
with DBSession() as session:
verify_query_engine_permission(engine_id, session=session)
uid = current_user.id
query_execution = logic.create_query_execution(query=query,
engine_id=engine_id,
uid=uid,
session=session)
data_doc = None
if data_cell_id:
datadoc_logic.append_query_executions_to_data_cell(
data_cell_id, [query_execution.id], session=session)
data_cell = datadoc_logic.get_data_cell_by_id(data_cell_id,
session=session)
data_doc = data_cell.doc
try:
run_query_task.apply_async(args=[
query_execution.id,
])
query_execution_dict = query_execution.to_dict()
if data_doc:
socketio.emit(
"data_doc_query_execution",
(
originator,
query_execution_dict,
data_cell_id,
),
namespace="/datadoc",
room=data_doc.id,
broadcast=True,
)
return query_execution_dict
except Exception as e:
# We might encounter ConnectionError caused by
# Redis connection failing
logic.create_query_execution_error(
query_execution.id,
error_type=None,
error_message_extracted="Encountered Error connecting to Redis",
error_message=str(e),
commit=False,
session=session,
)
query_execution.status = QueryExecutionStatus.ERROR
session.commit()
raise e
def get_query_engine_status(engine_id):
engine_checker = None
# Security check
with DBSession() as session:
verify_query_engine_permission(engine_id, session=session)
engine = admin_logic.get_query_engine_by_id(engine_id, session=session)
engine_checker = get_engine_checker_class(
engine.get_feature_params().get("status_checker", "NullChecker"))
api_assert(engine_checker is not None, "Invalid engine checker")
return engine_checker.check(engine_id=engine_id, uid=current_user.id)
def perform_table_upload():
import_config = json.loads(request.form["import_config"])
file_uploaded = request.files.get("file")
verify_import_config_permissions(import_config)
table_config = json.loads(request.form["table_config"])
engine_id = request.form["engine_id"]
verify_query_engine_permission(engine_id)
importer = get_importer(import_config, file_uploaded)
exporter = get_exporter(engine_id, current_user.id, table_config, importer)
return exporter.upload()
def update_query_snippet_by_id(
query_snippet_id,
context=None,
title=None,
description=None,
engine_id=None,
golden=None,
is_public=None,
):
updated_by = current_user.id
with DBSession() as session:
snippet = logic.get_snippet_by_id(query_snippet_id, session=session)
api_assert(snippet, "Snippet not found")
verify_query_engine_permission(snippet.engine_id, session=session)
# Pre condition check
if snippet.golden: # Editing a golden snippet
api_assert(
current_user.is_admin, "Golden snippet can only be updated by Admin"
)
elif not snippet.is_public: # Editing a private snippet
api_assert(
updated_by == snippet.created_by,
"Private query can only be updated by creator",
)
# Post condition check
if (
is_public is not None and is_public != snippet.is_public
): # Becomes public/private
api_assert(
updated_by == snippet.created_by,
"Only creator can toggle public/private",
)
if golden is not None and golden != snippet.golden: # Becomes golden
api_assert(current_user.is_admin, "Only data gurus can make snippet golden")
return logic.update_snippet_by_id(
query_snippet_id,
updated_by,
context=context,
title=title,
engine_id=engine_id,
description=description,
is_public=is_public,
golden=golden,
session=session,
)
def delete_query_snippet_by_id(query_snippet_id):
deleted_by = current_user.id
with DBSession() as session:
snippet = logic.get_snippet_by_id(query_snippet_id, session=session)
api_assert(snippet, "Snippet not found")
api_assert(
snippet.created_by == deleted_by or current_user.is_admin,
"Only creator or Admin can delete this snippet",
)
verify_query_engine_permission(snippet.engine_id, session=session)
logic.delete_snippet(query_snippet_id, deleted_by, session=session)
def cancel_query_execution(query_execution_id):
with DBSession() as session:
execution = logic.get_query_execution_by_id(query_execution_id, session=session)
verify_query_engine_permission(execution.engine_id, session=session)
execution_dict = execution.to_dict(True) if execution is not None else None
requestor = current_user.id
api_assert(
requestor == execution_dict["uid"], "You can only cancel your own queries"
)
if execution_dict and "task_id" in execution_dict:
task = run_query_task.AsyncResult(execution_dict["task_id"])
if task is not None:
task.abort()
def on_join_room(query_execution_id):
with DBSession() as session:
execution = qe_logic.get_query_execution_by_id(
query_execution_id, session=session
)
assert execution, "Invalid execution"
verify_query_engine_permission(execution.engine_id, session=session)
execution_dict = execution.to_dict(True) if execution is not None else None
join_room(query_execution_id)
if execution_dict and len(execution_dict.get("statement_executions", [])):
statement_execution = execution_dict["statement_executions"][-1]
# Format statement execution's logs
if statement_execution["has_log"]:
logs = qe_logic.get_statement_execution_stream_logs(
statement_execution["id"], from_end=True, session=session
)
statement_execution["log"] = [log.log for log in logs]
# Getting task's running data
if (
"task_id" in execution_dict
and execution_dict.get("status", None)
== QueryExecutionStatus.RUNNING.value
):
task = tasks.run_query_task.AsyncResult(execution_dict["task_id"])
try:
if task is not None and task.info is not None:
progress = task.info
if str(statement_execution["id"]) in progress:
statement_execution["percent_complete"] = progress[
str(statement_execution["id"])
].get("percent_complete")
execution_dict["total"] = progress.get("total", 0)
except Exception as e:
LOG.info(e)
emit(
"query",
execution_dict,
namespace=QUERY_EXECUTION_NAMESPACE,
room=query_execution_id,
)
def get_query_snippet_by_id(query_snippet_id):
with DBSession() as session:
query_snippet = logic.get_snippet_by_id(query_snippet_id, session=session)
verify_query_engine_permission(query_snippet.engine_id, session=session)
return query_snippet