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