def execute_sql_statements(
query_id,
rendered_query,
return_results=True,
store_results=False,
user_name=None,
session=None,
start_time=None,
expand_data=False,
): # pylint: disable=too-many-arguments, too-many-locals, too-many-statements
"""Executes the sql query returns the results."""
if store_results and start_time:
# only asynchronous queries
stats_logger.timing("sqllab.query.time_pending",
now_as_float() - start_time)
query = get_query(query_id, session)
payload = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
if database.allow_run_async and not results_backend:
raise SqlLabException("Results backend isn't configured.")
# Breaking down into multiple statements
parsed_query = ParsedQuery(rendered_query)
statements = parsed_query.get_statements()
logger.info(f"Query {query_id}: Executing {len(statements)} statement(s)")
logger.info(f"Query {query_id}: Set query to 'running'")
query.status = QueryStatus.RUNNING
query.start_running_time = now_as_float()
session.commit()
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=True,
user_name=user_name,
source=sources.get("sql_lab", None),
)
# Sharing a single connection and cursor across the
# execution of all statements (if many)
with closing(engine.raw_connection()) as conn:
with closing(conn.cursor()) as cursor:
statement_count = len(statements)
for i, statement in enumerate(statements):
# Check if stopped
query = get_query(query_id, session)
if query.status == QueryStatus.STOPPED:
return None
# Run statement
msg = f"Running statement {i+1} out of {statement_count}"
logger.info(f"Query {query_id}: {msg}")
query.set_extra_json_key("progress", msg)
session.commit()
try:
cdf = execute_sql_statement(statement, query, user_name,
session, cursor)
except Exception as e: # pylint: disable=broad-except
msg = str(e)
if statement_count > 1:
msg = f"[Statement {i+1} out of {statement_count}] " + msg
payload = handle_query_error(msg, query, session, payload)
return payload
# Success, updating the query entry in database
query.rows = cdf.size
query.progress = 100
query.set_extra_json_key("progress", None)
if query.select_as_cta:
query.select_sql = database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema,
show_cols=False,
latest_partition=False,
)
query.end_time = now_as_float()
data, selected_columns, all_columns, expanded_columns = _serialize_and_expand_data(
cdf, db_engine_spec, store_results and results_backend_use_msgpack,
expand_data)
payload.update({
"status": QueryStatus.SUCCESS,
"data": data,
"columns": all_columns,
"selected_columns": selected_columns,
"expanded_columns": expanded_columns,
"query": query.to_dict(),
})
payload["query"]["state"] = QueryStatus.SUCCESS
if store_results and results_backend:
key = str(uuid.uuid4())
logger.info(
f"Query {query_id}: Storing results in results backend, key: {key}"
)
with stats_timing("sqllab.query.results_backend_write", stats_logger):
with stats_timing(
"sqllab.query.results_backend_write_serialization",
stats_logger):
serialized_payload = _serialize_payload(
payload, results_backend_use_msgpack)
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config["CACHE_DEFAULT_TIMEOUT"]
compressed = zlib_compress(serialized_payload)
logger.debug(
f"*** serialized payload size: {getsizeof(serialized_payload)}"
)
logger.debug(
f"*** compressed payload size: {getsizeof(compressed)}")
results_backend.set(key, compressed, cache_timeout)
query.results_key = key
query.status = QueryStatus.SUCCESS
session.commit()
if return_results:
return payload
return None
def get_sql_results(self, query_id, return_results=True, store_results=False):
"""Executes the sql query returns the results."""
if not self.request.called_directly:
engine = sqlalchemy.create_engine(
app.config.get('SQLALCHEMY_DATABASE_URI'), poolclass=NullPool)
session_class = sessionmaker()
session_class.configure(bind=engine)
session = session_class()
else:
session = db.session()
session.commit() # HACK
try:
query = session.query(models.Query).filter_by(id=query_id).one()
except Exception as e:
logging.error("Query with id `{}` could not be retrieved".format(query_id))
logging.error("Sleeping for a sec and retrying...")
# Nasty hack to get around a race condition where the worker
# cannot find the query it's supposed to run
sleep(1)
query = session.query(models.Query).filter_by(id=query_id).one()
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
def handle_error(msg):
"""Local method handling error while processing the SQL"""
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
raise Exception(query.error_message)
if store_results and not results_backend:
handle_error("Results backend isn't configured.")
# Limit enforced only for retrieving the data, not for the CTA queries.
superset_query = SupersetQuery(query.sql)
executed_sql = superset_query.stripped()
if not superset_query.is_select() and not database.allow_dml:
handle_error(
"Only `SELECT` statements are allowed against this database")
if query.select_as_cta:
if not superset_query.is_select():
handle_error(
"Only `SELECT` statements can be used with the CREATE TABLE "
"feature.")
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = 'tmp_{}_table_{}'.format(
query.user_id,
start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
executed_sql = superset_query.as_create_table(query.tmp_table_name)
query.select_as_cta_used = True
elif (
query.limit and superset_query.is_select() and
db_engine_spec.limit_method == LimitMethod.WRAP_SQL):
executed_sql = database.wrap_sql_limit(executed_sql, query.limit)
query.limit_used = True
try:
template_processor = get_template_processor(
database=database, query=query)
executed_sql = template_processor.process_template(executed_sql)
executed_sql = db_engine_spec.sql_preprocessor(executed_sql)
except Exception as e:
logging.exception(e)
msg = "Template rendering failed: " + utils.error_msg_from_exception(e)
handle_error(msg)
query.executed_sql = executed_sql
query.status = QueryStatus.RUNNING
query.start_running_time = utils.now_as_float()
session.merge(query)
session.commit()
logging.info("Set query to 'running'")
engine = database.get_sqla_engine(schema=query.schema)
conn = engine.raw_connection()
cursor = conn.cursor()
logging.info("Running query: \n{}".format(executed_sql))
try:
logging.info(query.executed_sql)
cursor.execute(
query.executed_sql, **db_engine_spec.cursor_execute_kwargs)
except Exception as e:
logging.exception(e)
conn.close()
handle_error(db_engine_spec.extract_error_message(e))
try:
logging.info("Handling cursor")
db_engine_spec.handle_cursor(cursor, query, session)
logging.info("Fetching data: {}".format(query.to_dict()))
data = db_engine_spec.fetch_data(cursor, query.limit)
except Exception as e:
logging.exception(e)
conn.close()
handle_error(db_engine_spec.extract_error_message(e))
conn.commit()
conn.close()
if query.status == utils.QueryStatus.STOPPED:
return json.dumps({
'query_id': query.id,
'status': query.status,
'query': query.to_dict(),
}, default=utils.json_iso_dttm_ser)
column_names = (
[col[0] for col in cursor.description] if cursor.description else [])
column_names = dedup(column_names)
cdf = dataframe.SupersetDataFrame(pd.DataFrame(
list(data), columns=column_names))
query.rows = cdf.size
query.progress = 100
query.status = QueryStatus.SUCCESS
if query.select_as_cta:
query.select_sql = '{}'.format(database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema
))
query.end_time = utils.now_as_float()
session.merge(query)
session.flush()
payload = {
'query_id': query.id,
'status': query.status,
'data': cdf.data if cdf.data else [],
'columns': cdf.columns if cdf.columns else [],
'query': query.to_dict(),
}
payload = json.dumps(payload, default=utils.json_iso_dttm_ser)
if store_results:
key = '{}'.format(uuid.uuid4())
logging.info("Storing results in results backend, key: {}".format(key))
results_backend.set(key, zlib.compress(payload))
query.results_key = key
session.merge(query)
session.commit()
if return_results:
return payload
def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-locals, too-many-statements
query_id: int,
rendered_query: str,
return_results: bool,
store_results: bool,
user_name: Optional[str],
session: Session,
start_time: Optional[float],
expand_data: bool,
log_params: Optional[Dict[str, Any]],
) -> Optional[Dict[str, Any]]:
"""Executes the sql query returns the results."""
if store_results and start_time:
# only asynchronous queries
stats_logger.timing("sqllab.query.time_pending",
now_as_float() - start_time)
query = get_query(query_id, session)
payload: Dict[str, Any] = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
if database.allow_run_async and not results_backend:
raise SqlLabException("Results backend isn't configured.")
# Breaking down into multiple statements
parsed_query = ParsedQuery(rendered_query)
statements = parsed_query.get_statements()
logger.info("Query %s: Executing %i statement(s)", str(query_id),
len(statements))
logger.info("Query %s: Set query to 'running'", str(query_id))
query.status = QueryStatus.RUNNING
query.start_running_time = now_as_float()
session.commit()
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=True,
user_name=user_name,
source=QuerySource.SQL_LAB,
)
# Sharing a single connection and cursor across the
# execution of all statements (if many)
with closing(engine.raw_connection()) as conn:
with closing(conn.cursor()) as cursor:
statement_count = len(statements)
for i, statement in enumerate(statements):
# Check if stopped
query = get_query(query_id, session)
if query.status == QueryStatus.STOPPED:
return None
# Run statement
msg = f"Running statement {i+1} out of {statement_count}"
logger.info("Query %s: %s", str(query_id), msg)
query.set_extra_json_key("progress", msg)
session.commit()
try:
result_set = execute_sql_statement(statement, query,
user_name, session,
cursor, log_params)
except Exception as ex: # pylint: disable=broad-except
msg = str(ex)
if statement_count > 1:
msg = f"[Statement {i+1} out of {statement_count}] " + msg
payload = handle_query_error(msg, query, session, payload)
return payload
# Commit the connection so CTA queries will create the table.
conn.commit()
# Success, updating the query entry in database
query.rows = result_set.size
query.progress = 100
query.set_extra_json_key("progress", None)
if query.select_as_cta:
query.select_sql = database.select_star(
query.tmp_table_name,
schema=query.tmp_schema_name,
limit=query.limit,
show_cols=False,
latest_partition=False,
)
query.end_time = now_as_float()
use_arrow_data = store_results and cast(bool, results_backend_use_msgpack)
data, selected_columns, all_columns, expanded_columns = _serialize_and_expand_data(
result_set, db_engine_spec, use_arrow_data, expand_data)
# TODO: data should be saved separately from metadata (likely in Parquet)
payload.update({
"status": QueryStatus.SUCCESS,
"data": data,
"columns": all_columns,
"selected_columns": selected_columns,
"expanded_columns": expanded_columns,
"query": query.to_dict(),
})
payload["query"]["state"] = QueryStatus.SUCCESS
if store_results and results_backend:
key = str(uuid.uuid4())
logger.info("Query %s: Storing results in results backend, key: %s",
str(query_id), key)
with stats_timing("sqllab.query.results_backend_write", stats_logger):
with stats_timing(
"sqllab.query.results_backend_write_serialization",
stats_logger):
serialized_payload = _serialize_payload(
payload, cast(bool, results_backend_use_msgpack))
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config["CACHE_DEFAULT_TIMEOUT"]
compressed = zlib_compress(serialized_payload)
logger.debug("*** serialized payload size: %i",
getsizeof(serialized_payload))
logger.debug("*** compressed payload size: %i",
getsizeof(compressed))
results_backend.set(key, compressed, cache_timeout)
query.results_key = key
query.status = QueryStatus.SUCCESS
session.commit()
if return_results:
# since we're returning results we need to create non-arrow data
if use_arrow_data:
(
data,
selected_columns,
all_columns,
expanded_columns,
) = _serialize_and_expand_data(result_set, db_engine_spec, False,
expand_data)
payload.update({
"data": data,
"columns": all_columns,
"selected_columns": selected_columns,
"expanded_columns": expanded_columns,
})
return payload
return None
def execute_sql(
ctask, query_id, rendered_query, return_results=True, store_results=False,
user_name=None, session=None, start_time=None,
):
"""Executes the sql query returns the results."""
if store_results and start_time:
# only asynchronous queries
stats_logger.timing(
'sqllab.query.time_pending', now_as_float() - start_time)
query = get_query(query_id, session)
payload = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
def handle_error(msg):
"""Local method handling error while processing the SQL"""
troubleshooting_link = config['TROUBLESHOOTING_LINK']
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
payload.update({
'status': query.status,
'error': msg,
})
if troubleshooting_link:
payload['link'] = troubleshooting_link
return payload
if store_results and not results_backend:
return handle_error("Results backend isn't configured.")
# Limit enforced only for retrieving the data, not for the CTA queries.
superset_query = SupersetQuery(rendered_query)
executed_sql = superset_query.stripped()
SQL_MAX_ROWS = app.config.get('SQL_MAX_ROW')
if not superset_query.is_readonly() and not database.allow_dml:
return handle_error(
'Only `SELECT` statements are allowed against this database')
if query.select_as_cta:
if not superset_query.is_select():
return handle_error(
'Only `SELECT` statements can be used with the CREATE TABLE '
'feature.')
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = 'tmp_{}_table_{}'.format(
query.user_id, start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
executed_sql = superset_query.as_create_table(query.tmp_table_name)
query.select_as_cta_used = True
if (superset_query.is_select() and SQL_MAX_ROWS and
(not query.limit or query.limit > SQL_MAX_ROWS)):
query.limit = SQL_MAX_ROWS
executed_sql = database.apply_limit_to_sql(executed_sql, query.limit)
# Hook to allow environment-specific mutation (usually comments) to the SQL
SQL_QUERY_MUTATOR = config.get('SQL_QUERY_MUTATOR')
if SQL_QUERY_MUTATOR:
executed_sql = SQL_QUERY_MUTATOR(
executed_sql, user_name, security_manager, database)
query.executed_sql = executed_sql
query.status = QueryStatus.RUNNING
query.start_running_time = now_as_float()
session.merge(query)
session.commit()
logging.info("Set query to 'running'")
conn = None
try:
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=True,
user_name=user_name,
)
conn = engine.raw_connection()
cursor = conn.cursor()
logging.info('Running query: \n{}'.format(executed_sql))
logging.info(query.executed_sql)
query_start_time = now_as_float()
db_engine_spec.execute(cursor, query.executed_sql, async_=True)
logging.info('Handling cursor')
db_engine_spec.handle_cursor(cursor, query, session)
logging.info('Fetching data: {}'.format(query.to_dict()))
stats_logger.timing(
'sqllab.query.time_executing_query',
now_as_float() - query_start_time)
fetching_start_time = now_as_float()
data = db_engine_spec.fetch_data(cursor, query.limit)
stats_logger.timing(
'sqllab.query.time_fetching_results',
now_as_float() - fetching_start_time)
except SoftTimeLimitExceeded as e:
logging.exception(e)
if conn is not None:
conn.close()
return handle_error(
"SQL Lab timeout. This environment's policy is to kill queries "
'after {} seconds.'.format(SQLLAB_TIMEOUT))
except Exception as e:
logging.exception(e)
if conn is not None:
conn.close()
return handle_error(db_engine_spec.extract_error_message(e))
logging.info('Fetching cursor description')
cursor_description = cursor.description
if conn is not None:
conn.commit()
conn.close()
if query.status == QueryStatus.STOPPED:
return handle_error('The query has been stopped')
cdf = dataframe.SupersetDataFrame(data, cursor_description, db_engine_spec)
query.rows = cdf.size
query.progress = 100
query.status = QueryStatus.SUCCESS
if query.select_as_cta:
query.select_sql = '{}'.format(
database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema,
show_cols=False,
latest_partition=False))
query.end_time = now_as_float()
session.merge(query)
session.flush()
payload.update({
'status': query.status,
'data': cdf.data if cdf.data else [],
'columns': cdf.columns if cdf.columns else [],
'query': query.to_dict(),
})
if store_results:
key = '{}'.format(uuid.uuid4())
logging.info('Storing results in results backend, key: {}'.format(key))
write_to_results_backend_start = now_as_float()
json_payload = json.dumps(
payload, default=json_iso_dttm_ser, ignore_nan=True)
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config.get('CACHE_DEFAULT_TIMEOUT', 0)
results_backend.set(key, zlib_compress(json_payload), cache_timeout)
query.results_key = key
stats_logger.timing(
'sqllab.query.results_backend_write',
now_as_float() - write_to_results_backend_start)
session.merge(query)
session.commit()
if return_results:
return payload
def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-locals, too-many-statements, too-many-branches
query_id: int,
rendered_query: str,
return_results: bool,
store_results: bool,
user_name: Optional[str],
session: Session,
start_time: Optional[float],
expand_data: bool,
log_params: Optional[Dict[str, Any]],
) -> Optional[Dict[str, Any]]:
"""Executes the sql query returns the results."""
if store_results and start_time:
# only asynchronous queries
stats_logger.timing("sqllab.query.time_pending", now_as_float() - start_time)
query = get_query(query_id, session)
payload: Dict[str, Any] = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
if database.allow_run_async and not results_backend:
raise SupersetErrorException(
SupersetError(
message=__("Results backend is not configured."),
error_type=SupersetErrorType.RESULTS_BACKEND_NOT_CONFIGURED_ERROR,
level=ErrorLevel.ERROR,
)
)
# Breaking down into multiple statements
parsed_query = ParsedQuery(rendered_query, strip_comments=True)
if not db_engine_spec.run_multiple_statements_as_one:
statements = parsed_query.get_statements()
logger.info(
"Query %s: Executing %i statement(s)", str(query_id), len(statements)
)
else:
statements = [rendered_query]
logger.info("Query %s: Executing query as a single statement", str(query_id))
logger.info("Query %s: Set query to 'running'", str(query_id))
query.status = QueryStatus.RUNNING
query.start_running_time = now_as_float()
session.commit()
# Should we create a table or view from the select?
if (
query.select_as_cta
and query.ctas_method == CtasMethod.TABLE
and not parsed_query.is_valid_ctas()
):
raise SupersetErrorException(
SupersetError(
message=__(
"CTAS (create table as select) can only be run with a query where "
"the last statement is a SELECT. Please make sure your query has "
"a SELECT as its last statement. Then, try running your query "
"again."
),
error_type=SupersetErrorType.INVALID_CTAS_QUERY_ERROR,
level=ErrorLevel.ERROR,
)
)
if (
query.select_as_cta
and query.ctas_method == CtasMethod.VIEW
and not parsed_query.is_valid_cvas()
):
raise SupersetErrorException(
SupersetError(
message=__(
"CVAS (create view as select) can only be run with a query with "
"a single SELECT statement. Please make sure your query has only "
"a SELECT statement. Then, try running your query again."
),
error_type=SupersetErrorType.INVALID_CVAS_QUERY_ERROR,
level=ErrorLevel.ERROR,
)
)
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=True,
user_name=user_name,
source=QuerySource.SQL_LAB,
)
# Sharing a single connection and cursor across the
# execution of all statements (if many)
with closing(engine.raw_connection()) as conn:
# closing the connection closes the cursor as well
cursor = conn.cursor()
statement_count = len(statements)
for i, statement in enumerate(statements):
# Check if stopped
query = get_query(query_id, session)
if query.status == QueryStatus.STOPPED:
return None
# For CTAS we create the table only on the last statement
apply_ctas = query.select_as_cta and (
query.ctas_method == CtasMethod.VIEW
or (query.ctas_method == CtasMethod.TABLE and i == len(statements) - 1)
)
# Run statement
msg = f"Running statement {i+1} out of {statement_count}"
logger.info("Query %s: %s", str(query_id), msg)
query.set_extra_json_key("progress", msg)
session.commit()
try:
result_set = execute_sql_statement(
statement,
query,
user_name,
session,
cursor,
log_params,
apply_ctas,
)
except Exception as ex: # pylint: disable=broad-except
msg = str(ex)
prefix_message = (
f"[Statement {i+1} out of {statement_count}]"
if statement_count > 1
else ""
)
payload = handle_query_error(
ex, query, session, payload, prefix_message
)
return payload
# Commit the connection so CTA queries will create the table.
conn.commit()
# Success, updating the query entry in database
query.rows = result_set.size
query.progress = 100
query.set_extra_json_key("progress", None)
if query.select_as_cta:
query.select_sql = database.select_star(
query.tmp_table_name,
schema=query.tmp_schema_name,
limit=query.limit,
show_cols=False,
latest_partition=False,
)
query.end_time = now_as_float()
use_arrow_data = store_results and cast(bool, results_backend_use_msgpack)
data, selected_columns, all_columns, expanded_columns = _serialize_and_expand_data(
result_set, db_engine_spec, use_arrow_data, expand_data
)
# TODO: data should be saved separately from metadata (likely in Parquet)
payload.update(
{
"status": QueryStatus.SUCCESS,
"data": data,
"columns": all_columns,
"selected_columns": selected_columns,
"expanded_columns": expanded_columns,
"query": query.to_dict(),
}
)
payload["query"]["state"] = QueryStatus.SUCCESS
if store_results and results_backend:
key = str(uuid.uuid4())
logger.info(
"Query %s: Storing results in results backend, key: %s", str(query_id), key
)
with stats_timing("sqllab.query.results_backend_write", stats_logger):
with stats_timing(
"sqllab.query.results_backend_write_serialization", stats_logger
):
serialized_payload = _serialize_payload(
payload, cast(bool, results_backend_use_msgpack)
)
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config["CACHE_DEFAULT_TIMEOUT"]
compressed = zlib_compress(serialized_payload)
logger.debug(
"*** serialized payload size: %i", getsizeof(serialized_payload)
)
logger.debug("*** compressed payload size: %i", getsizeof(compressed))
results_backend.set(key, compressed, cache_timeout)
query.results_key = key
query.status = QueryStatus.SUCCESS
session.commit()
if return_results:
# since we're returning results we need to create non-arrow data
if use_arrow_data:
(
data,
selected_columns,
all_columns,
expanded_columns,
) = _serialize_and_expand_data(
result_set, db_engine_spec, False, expand_data
)
payload.update(
{
"data": data,
"columns": all_columns,
"selected_columns": selected_columns,
"expanded_columns": expanded_columns,
}
)
return payload
return None
def get_sql_results(self, query_id, return_results=True, store_results=False):
"""Executes the sql query returns the results."""
if not self.request.called_directly:
engine = sqlalchemy.create_engine(
app.config.get('SQLALCHEMY_DATABASE_URI'), poolclass=NullPool)
session_class = sessionmaker()
session_class.configure(bind=engine)
session = session_class()
else:
session = db.session()
session.commit() # HACK
query = session.query(models.Query).filter_by(id=query_id).one()
database = query.database
db_engine_spec = database.db_engine_spec
def handle_error(msg):
"""Local method handling error while processing the SQL"""
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
raise Exception(query.error_message)
if store_results and not results_backend:
handle_error("Results backend isn't configured.")
# Limit enforced only for retrieving the data, not for the CTA queries.
superset_query = SupersetQuery(query.sql)
executed_sql = superset_query.stripped()
if not superset_query.is_select() and not database.allow_dml:
handle_error(
"Only `SELECT` statements are allowed against this database")
if query.select_as_cta:
if not superset_query.is_select():
handle_error(
"Only `SELECT` statements can be used with the CREATE TABLE "
"feature.")
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = 'tmp_{}_table_{}'.format(
query.user_id,
start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
executed_sql = superset_query.as_create_table(query.tmp_table_name)
query.select_as_cta_used = True
elif (
query.limit and superset_query.is_select() and
db_engine_spec.limit_method == LimitMethod.WRAP_SQL):
executed_sql = database.wrap_sql_limit(executed_sql, query.limit)
query.limit_used = True
engine = database.get_sqla_engine(schema=query.schema)
try:
template_processor = get_template_processor(
database=database, query=query)
executed_sql = template_processor.process_template(executed_sql)
executed_sql = db_engine_spec.sql_preprocessor(executed_sql)
except Exception as e:
logging.exception(e)
msg = "Template rendering failed: " + utils.error_msg_from_exception(e)
handle_error(msg)
query.executed_sql = executed_sql
logging.info("Running query: \n{}".format(executed_sql))
try:
result_proxy = engine.execute(query.executed_sql, schema=query.schema)
except Exception as e:
logging.exception(e)
handle_error(db_engine_spec.extract_error_message(e))
cursor = result_proxy.cursor
query.status = QueryStatus.RUNNING
session.flush()
db_engine_spec.handle_cursor(cursor, query, session)
cdf = None
if result_proxy.cursor:
column_names = [col[0] for col in result_proxy.cursor.description]
column_names = dedup(column_names)
if db_engine_spec.limit_method == LimitMethod.FETCH_MANY:
data = result_proxy.fetchmany(query.limit)
else:
data = result_proxy.fetchall()
cdf = dataframe.SupersetDataFrame(
pd.DataFrame(data, columns=column_names))
query.rows = result_proxy.rowcount
query.progress = 100
query.status = QueryStatus.SUCCESS
if query.rows == -1 and cdf:
# Presto doesn't provide result_proxy.row_count
query.rows = cdf.size
if query.select_as_cta:
query.select_sql = '{}'.format(database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema
))
query.end_time = utils.now_as_float()
session.flush()
payload = {
'query_id': query.id,
'status': query.status,
'data': [],
}
payload['data'] = cdf.data if cdf else []
payload['columns'] = cdf.columns_dict if cdf else []
payload['query'] = query.to_dict()
payload = json.dumps(payload, default=utils.json_iso_dttm_ser)
if store_results:
key = '{}'.format(uuid.uuid4())
logging.info("Storing results in results backend, key: {}".format(key))
results_backend.set(key, zlib.compress(payload))
query.results_key = key
session.flush()
session.commit()
if return_results:
return payload
def execute_sql(
ctask,
query_id,
return_results=True,
store_results=False,
user_name=None,
template_params=None,
):
"""Executes the sql query returns the results."""
session = get_session(not ctask.request.called_directly)
query = get_query(query_id, session)
payload = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
def handle_error(msg):
"""Local method handling error while processing the SQL"""
troubleshooting_link = config['TROUBLESHOOTING_LINK']
msg = 'Error: {}. You can find common superset errors and their \
resolutions at: {}' .format(msg, troubleshooting_link) \
if troubleshooting_link else msg
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
payload.update({
'status': query.status,
'error': msg,
})
return payload
if store_results and not results_backend:
return handle_error("Results backend isn't configured.")
try:
template_processor = get_template_processor(database=database,
query=query)
tp = template_params or {}
rendered_query = template_processor.process_template(query.sql, **tp)
except Exception as e:
logging.exception(e)
msg = 'Template rendering failed: ' + utils.error_msg_from_exception(e)
return handle_error(msg)
# Limit enforced only for retrieving the data, not for the CTA queries.
superset_query = SupersetQuery(rendered_query)
executed_sql = superset_query.stripped()
if not superset_query.is_select() and not database.allow_dml:
return handle_error(
'Only `SELECT` statements are allowed against this database')
if query.select_as_cta:
if not superset_query.is_select():
return handle_error(
'Only `SELECT` statements can be used with the CREATE TABLE '
'feature.')
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = 'tmp_{}_table_{}'.format(
query.user_id, start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
executed_sql = superset_query.as_create_table(query.tmp_table_name)
query.select_as_cta_used = True
elif (query.limit and superset_query.is_select()
and db_engine_spec.limit_method == LimitMethod.WRAP_SQL):
executed_sql = database.wrap_sql_limit(executed_sql, query.limit)
query.limit_used = True
# Hook to allow environment-specific mutation (usually comments) to the SQL
SQL_QUERY_MUTATOR = config.get('SQL_QUERY_MUTATOR')
if SQL_QUERY_MUTATOR:
executed_sql = SQL_QUERY_MUTATOR(executed_sql, user_name,
security_manager, database)
query.executed_sql = executed_sql
query.status = QueryStatus.RUNNING
query.start_running_time = utils.now_as_float()
session.merge(query)
session.commit()
logging.info("Set query to 'running'")
conn = None
try:
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=not ctask.request.called_directly,
user_name=user_name,
)
conn = engine.raw_connection()
cursor = conn.cursor()
logging.info('Running query: \n{}'.format(executed_sql))
logging.info(query.executed_sql)
cursor.execute(query.executed_sql,
**db_engine_spec.cursor_execute_kwargs)
logging.info('Handling cursor')
db_engine_spec.handle_cursor(cursor, query, session)
logging.info('Fetching data: {}'.format(query.to_dict()))
data = db_engine_spec.fetch_data(cursor, query.limit)
except SoftTimeLimitExceeded as e:
logging.exception(e)
if conn is not None:
conn.close()
return handle_error(
"SQL Lab timeout. This environment's policy is to kill queries "
'after {} seconds.'.format(SQLLAB_TIMEOUT))
except Exception as e:
logging.exception(e)
if conn is not None:
conn.close()
return handle_error(db_engine_spec.extract_error_message(e))
logging.info('Fetching cursor description')
cursor_description = cursor.description
if conn is not None:
conn.commit()
conn.close()
if query.status == utils.QueryStatus.STOPPED:
return json.dumps(
{
'query_id': query.id,
'status': query.status,
'query': query.to_dict(),
},
default=utils.json_iso_dttm_ser)
cdf = convert_results_to_df(cursor_description, data)
query.rows = cdf.size
query.progress = 100
query.status = QueryStatus.SUCCESS
if query.select_as_cta:
query.select_sql = '{}'.format(
database.select_star(query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema,
show_cols=False,
latest_partition=False))
query.end_time = utils.now_as_float()
session.merge(query)
session.flush()
payload.update({
'status': query.status,
'data': cdf.data if cdf.data else [],
'columns': cdf.columns if cdf.columns else [],
'query': query.to_dict(),
})
if store_results:
key = '{}'.format(uuid.uuid4())
logging.info('Storing results in results backend, key: {}'.format(key))
json_payload = json.dumps(payload, default=utils.json_iso_dttm_ser)
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config.get('CACHE_DEFAULT_TIMEOUT', 0)
results_backend.set(key, utils.zlib_compress(json_payload),
cache_timeout)
query.results_key = key
query.end_result_backend_time = utils.now_as_float()
session.merge(query)
session.commit()
if return_results:
return payload
def execute_sql(ctask,
query_id,
return_results=True,
store_results=False,
user_name=None):
"""Executes the sql query returns the results."""
session = get_session(not ctask.request.called_directly)
query = get_query(query_id, session)
payload = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
def handle_error(msg):
"""Local method handling error while processing the SQL"""
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
payload.update({
'status': query.status,
'error': msg,
})
return payload
if store_results and not results_backend:
return handle_error("Results backend isn't configured.")
# Limit enforced only for retrieving the data, not for the CTA queries.
superset_query = SupersetQuery(query.sql)
executed_sql = superset_query.stripped()
if not superset_query.is_select() and not database.allow_dml:
return handle_error(
"Only `SELECT` statements are allowed against this database")
if query.select_as_cta:
if not superset_query.is_select():
return handle_error(
"Only `SELECT` statements can be used with the CREATE TABLE "
"feature.")
return
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = 'tmp_{}_table_{}'.format(
query.user_id, start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
executed_sql = superset_query.as_create_table(query.tmp_table_name)
query.select_as_cta_used = True
elif (query.limit and superset_query.is_select()
and db_engine_spec.limit_method == LimitMethod.WRAP_SQL):
executed_sql = database.wrap_sql_limit(executed_sql, query.limit)
query.limit_used = True
try:
template_processor = get_template_processor(database=database,
query=query)
executed_sql = template_processor.process_template(executed_sql)
except Exception as e:
logging.exception(e)
msg = "Template rendering failed: " + utils.error_msg_from_exception(e)
return handle_error(msg)
query.executed_sql = executed_sql
query.status = QueryStatus.RUNNING
query.start_running_time = utils.now_as_float()
session.merge(query)
session.commit()
logging.info("Set query to 'running'")
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=not ctask.request.called_directly,
user_name=user_name)
try:
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=not ctask.request.called_directly,
user_name=user_name)
conn = engine.raw_connection()
cursor = conn.cursor()
logging.info("Running query: \n{}".format(executed_sql))
logging.info(query.executed_sql)
cursor.execute(query.executed_sql,
**db_engine_spec.cursor_execute_kwargs)
logging.info("Handling cursor")
db_engine_spec.handle_cursor(cursor, query, session)
logging.info("Fetching data: {}".format(query.to_dict()))
data = db_engine_spec.fetch_data(cursor, query.limit)
except SoftTimeLimitExceeded as e:
logging.exception(e)
conn.close()
return handle_error(
"SQL Lab timeout. This environment's policy is to kill queries "
"after {} seconds.".format(SQLLAB_TIMEOUT))
except Exception as e:
logging.exception(e)
conn.close()
return handle_error(db_engine_spec.extract_error_message(e))
logging.info("Fetching cursor description")
cursor_description = cursor.description
conn.commit()
conn.close()
if query.status == utils.QueryStatus.STOPPED:
return json.dumps(
{
'query_id': query.id,
'status': query.status,
'query': query.to_dict(),
},
default=utils.json_iso_dttm_ser)
column_names = ([col[0] for col in cursor_description]
if cursor_description else [])
column_names = dedup(column_names)
cdf = dataframe.SupersetDataFrame(
pd.DataFrame(list(data), columns=column_names))
query.rows = cdf.size
query.progress = 100
query.status = QueryStatus.SUCCESS
if query.select_as_cta:
query.select_sql = '{}'.format(
database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema,
show_cols=False,
latest_partition=False,
))
query.end_time = utils.now_as_float()
session.merge(query)
session.flush()
payload.update({
'status': query.status,
'data': cdf.data if cdf.data else [],
'columns': cdf.columns if cdf.columns else [],
'query': query.to_dict(),
})
if store_results:
key = '{}'.format(uuid.uuid4())
logging.info("Storing results in results backend, key: {}".format(key))
json_payload = json.dumps(payload, default=utils.json_iso_dttm_ser)
results_backend.set(key, utils.zlib_compress(json_payload))
query.results_key = key
query.end_result_backend_time = utils.now_as_float()
session.merge(query)
session.commit()
if return_results:
return payload
def execute_norm(ctask,
query_id,
rendered_query,
return_results=True,
store_results=False,
user_name=None,
session=None):
""" Executes the norm script and returns the results"""
if rendered_query.lower().find('select') >= 0:
return execute_sql(ctask, query_id, rendered_query, return_results,
store_results, user_name, session)
query = get_query(query_id, session)
payload = dict(query_id=query_id)
def handle_error(msg):
"""Local method handling error while processing the SQL"""
troubleshooting_link = config['TROUBLESHOOTING_LINK']
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
payload.update({
'status': query.status,
'error': msg,
})
if troubleshooting_link:
payload['link'] = troubleshooting_link
return payload
if store_results and not results_backend:
return handle_error("Results backend isn't configured.")
query.executed_sql = rendered_query
query.status = QueryStatus.RUNNING
query.start_running_time = utils.now_as_float()
session.merge(query)
session.commit()
logging.info("Set query to 'running'")
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
try:
data = norm.execute(query.executed_sql, session, query.user)
except SoftTimeLimitExceeded as e:
logging.exception(e)
return handle_error(
"SQL Lab timeout. This environment's policy is to kill queries "
'after {} seconds.'.format(SQLLAB_TIMEOUT))
except Exception as e:
logging.exception(e)
return handle_error(db_engine_spec.extract_error_message(e))
if query.status == utils.QueryStatus.STOPPED:
return handle_error('The query has been stopped')
cdf = dataframe.SupersetDataFrame(data)
query.rows = cdf.size
query.progress = 100
query.status = QueryStatus.SUCCESS
query.end_time = utils.now_as_float()
session.merge(query)
session.flush()
payload.update({
'status': query.status,
'data': cdf.data if cdf.data else [],
'columns': cdf.columns if cdf.columns else [],
'query': query.to_dict(),
})
if store_results:
key = '{}'.format(uuid.uuid4())
logging.info('Storing results in results backend, key: {}'.format(key))
json_payload = json.dumps(payload, default=utils.json_iso_dttm_ser)
cache_timeout = config.get('CACHE_DEFAULT_TIMEOUT', 0)
results_backend.set(key, utils.zlib_compress(json_payload),
cache_timeout)
query.results_key = key
query.end_result_backend_time = utils.now_as_float()
session.merge(query)
session.commit()
if return_results:
return payload
def execute_sql_statements(
ctask,
query_id,
rendered_query,
return_results=True,
store_results=False,
user_name=None,
session=None,
start_time=None,
):
"""Executes the sql query returns the results."""
if store_results and start_time:
# only asynchronous queries
stats_logger.timing("sqllab.query.time_pending", now_as_float() - start_time)
query = get_query(query_id, session)
payload = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
if store_results and not results_backend:
raise SqlLabException("Results backend isn't configured.")
# Breaking down into multiple statements
parsed_query = ParsedQuery(rendered_query)
statements = parsed_query.get_statements()
logging.info(f"Executing {len(statements)} statement(s)")
logging.info("Set query to 'running'")
query.status = QueryStatus.RUNNING
query.start_running_time = now_as_float()
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=True,
user_name=user_name,
source=sources.get("sql_lab", None),
)
# Sharing a single connection and cursor across the
# execution of all statements (if many)
with closing(engine.raw_connection()) as conn:
with closing(conn.cursor()) as cursor:
query.connection_id = db_engine_spec.get_connection_id(cursor)
session.commit()
statement_count = len(statements)
for i, statement in enumerate(statements):
# check if the query was stopped
session.refresh(query)
if query.status == QueryStatus.STOPPED:
payload.update({"status": query.status})
return payload
msg = f"Running statement {i+1} out of {statement_count}"
logging.info(msg)
query.set_extra_json_key("progress", msg)
session.commit()
try:
cdf = execute_sql_statement(
statement, query, user_name, session, cursor
)
msg = f"Running statement {i+1} out of {statement_count}"
except Exception as e:
# query can be stopped in another thread/worker
# but in synchronized mode it may lead to an error
# skip error the error in such case
session.refresh(query)
if query.status == QueryStatus.STOPPED:
payload.update({"status": query.status})
return payload
msg = str(e)
if statement_count > 1:
msg = f"[Statement {i+1} out of {statement_count}] " + msg
payload = handle_query_error(msg, query, session, payload)
return payload
# Success, updating the query entry in database
query.rows = cdf.size
query.progress = 100
query.set_extra_json_key("progress", None)
query.connection_id = None
if query.select_as_cta:
query.select_sql = database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema,
show_cols=False,
latest_partition=False,
)
query.end_time = now_as_float()
selected_columns = cdf.columns or []
data = cdf.data or []
all_columns, data, expanded_columns = db_engine_spec.expand_data(
selected_columns, data
)
payload.update(
{
"status": QueryStatus.SUCCESS,
"data": data,
"columns": all_columns,
"selected_columns": selected_columns,
"expanded_columns": expanded_columns,
"query": query.to_dict(),
}
)
payload["query"]["state"] = QueryStatus.SUCCESS
# go over each row, find bytes columns that start with the magic UAST
# sequence b'\x00bgr', and replace it with a string containing the
# UAST in JSON
for row in payload["data"]:
for k, v in row.items():
if isinstance(v, bytes) and len(v) > 4 and v[0:4] == b"\x00bgr":
try:
ctx = decode(v, format=0)
row[k] = json.dumps(ctx.load())
except Exception:
pass
if store_results:
key = str(uuid.uuid4())
logging.info(f"Storing results in results backend, key: {key}")
with stats_timing("sqllab.query.results_backend_write", stats_logger):
json_payload = json.dumps(
payload, default=json_iso_dttm_ser, ignore_nan=True
)
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config.get("CACHE_DEFAULT_TIMEOUT", 0)
results_backend.set(key, zlib_compress(json_payload), cache_timeout)
query.results_key = key
query.status = QueryStatus.SUCCESS
session.commit()
if return_results:
return payload
def execute_sql(
ctask, query_id, rendered_query, return_results=True, store_results=False,
user_name=None, session=None,
):
"""Executes the sql query returns the results."""
query = get_query(query_id, session)
payload = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
def handle_error(msg):
"""Local method handling error while processing the SQL"""
troubleshooting_link = config['TROUBLESHOOTING_LINK']
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
payload.update({
'status': query.status,
'error': msg,
})
if troubleshooting_link:
payload['link'] = troubleshooting_link
return payload
if store_results and not results_backend:
return handle_error("Results backend isn't configured.")
# Limit enforced only for retrieving the data, not for the CTA queries.
superset_query = SupersetQuery(rendered_query)
executed_sql = superset_query.stripped()
SQL_MAX_ROWS = app.config.get('SQL_MAX_ROW')
if not superset_query.is_select() and not database.allow_dml:
return handle_error(
'Only `SELECT` statements are allowed against this database')
if query.select_as_cta:
if not superset_query.is_select():
return handle_error(
'Only `SELECT` statements can be used with the CREATE TABLE '
'feature.')
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = 'tmp_{}_table_{}'.format(
query.user_id, start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
executed_sql = superset_query.as_create_table(query.tmp_table_name)
query.select_as_cta_used = True
if (superset_query.is_select() and SQL_MAX_ROWS and
(not query.limit or query.limit > SQL_MAX_ROWS)):
query.limit = SQL_MAX_ROWS
executed_sql = database.apply_limit_to_sql(executed_sql, query.limit)
# Hook to allow environment-specific mutation (usually comments) to the SQL
SQL_QUERY_MUTATOR = config.get('SQL_QUERY_MUTATOR')
if SQL_QUERY_MUTATOR:
executed_sql = SQL_QUERY_MUTATOR(
executed_sql, user_name, security_manager, database)
query.executed_sql = executed_sql
query.status = QueryStatus.RUNNING
query.start_running_time = utils.now_as_float()
session.merge(query)
session.commit()
logging.info("Set query to 'running'")
conn = None
try:
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=not ctask.request.called_directly,
user_name=user_name,
)
conn = engine.raw_connection()
cursor = conn.cursor()
logging.info('Running query: \n{}'.format(executed_sql))
logging.info(query.executed_sql)
cursor.execute(query.executed_sql,
**db_engine_spec.cursor_execute_kwargs)
logging.info('Handling cursor')
db_engine_spec.handle_cursor(cursor, query, session)
logging.info('Fetching data: {}'.format(query.to_dict()))
data = db_engine_spec.fetch_data(cursor, query.limit)
except SoftTimeLimitExceeded as e:
logging.exception(e)
if conn is not None:
conn.close()
return handle_error(
"SQL Lab timeout. This environment's policy is to kill queries "
'after {} seconds.'.format(SQLLAB_TIMEOUT))
except Exception as e:
logging.exception(e)
if conn is not None:
conn.close()
return handle_error(db_engine_spec.extract_error_message(e))
logging.info('Fetching cursor description')
column_names = db_engine_spec.get_normalized_column_names(cursor.description)
if conn is not None:
conn.commit()
conn.close()
if query.status == utils.QueryStatus.STOPPED:
return handle_error('The query has been stopped')
cdf = convert_results_to_df(column_names, data)
query.rows = cdf.size
query.progress = 100
query.status = QueryStatus.SUCCESS
if query.select_as_cta:
query.select_sql = '{}'.format(
database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema,
show_cols=False,
latest_partition=False))
query.end_time = utils.now_as_float()
session.merge(query)
session.flush()
payload.update({
'status': query.status,
'data': cdf.data if cdf.data else [],
'columns': cdf.columns if cdf.columns else [],
'query': query.to_dict(),
})
if store_results:
key = '{}'.format(uuid.uuid4())
logging.info('Storing results in results backend, key: {}'.format(key))
json_payload = json.dumps(payload, default=utils.json_iso_dttm_ser)
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config.get('CACHE_DEFAULT_TIMEOUT', 0)
results_backend.set(key, utils.zlib_compress(json_payload), cache_timeout)
query.results_key = key
query.end_result_backend_time = utils.now_as_float()
session.merge(query)
session.commit()
if return_results:
return payload
def execute_sql(
ctask, query_id, return_results=True, store_results=False, user_name=None,
template_params=None,
):
"""Executes the sql query returns the results."""
session = get_session(not ctask.request.called_directly)
query = get_query(query_id, session)
payload = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
def handle_error(msg):
"""Local method handling error while processing the SQL"""
troubleshooting_link = config['TROUBLESHOOTING_LINK']
msg = 'Error: {}. You can find common superset errors and their \
resolutions at: {}'.format(msg, troubleshooting_link) \
if troubleshooting_link else msg
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
payload.update({
'status': query.status,
'error': msg,
})
return payload
if store_results and not results_backend:
return handle_error("Results backend isn't configured.")
# Limit enforced only for retrieving the data, not for the CTA queries.
superset_query = SupersetQuery(query.sql)
executed_sql = superset_query.stripped()
if not superset_query.is_select() and not database.allow_dml:
return handle_error(
'Only `SELECT` statements are allowed against this database')
if query.select_as_cta:
if not superset_query.is_select():
return handle_error(
'Only `SELECT` statements can be used with the CREATE TABLE '
'feature.')
return
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = 'tmp_{}_table_{}'.format(
query.user_id, start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
executed_sql = superset_query.as_create_table(query.tmp_table_name)
query.select_as_cta_used = True
elif (query.limit and superset_query.is_select() and
db_engine_spec.limit_method == LimitMethod.WRAP_SQL):
executed_sql = database.wrap_sql_limit(executed_sql, query.limit)
query.limit_used = True
try:
template_processor = get_template_processor(
database=database, query=query)
tp = template_params or {}
executed_sql = template_processor.process_template(
executed_sql, **tp)
except Exception as e:
logging.exception(e)
msg = 'Template rendering failed: ' + utils.error_msg_from_exception(e)
return handle_error(msg)
query.executed_sql = executed_sql
query.status = QueryStatus.RUNNING
query.start_running_time = utils.now_as_float()
session.merge(query)
session.commit()
logging.info("Set query to 'running'")
conn = None
try:
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=not ctask.request.called_directly,
user_name=user_name,
)
conn = engine.raw_connection()
cursor = conn.cursor()
logging.info('Running query: \n{}'.format(executed_sql))
logging.info(query.executed_sql)
cursor.execute(query.executed_sql,
**db_engine_spec.cursor_execute_kwargs)
logging.info('Handling cursor')
db_engine_spec.handle_cursor(cursor, query, session)
logging.info('Fetching data: {}'.format(query.to_dict()))
data = db_engine_spec.fetch_data(cursor, query.limit)
except SoftTimeLimitExceeded as e:
logging.exception(e)
if conn is not None:
conn.close()
return handle_error(
"SQL Lab timeout. This environment's policy is to kill queries "
'after {} seconds.'.format(SQLLAB_TIMEOUT))
except Exception as e:
logging.exception(e)
if conn is not None:
conn.close()
return handle_error(db_engine_spec.extract_error_message(e))
logging.info('Fetching cursor description')
cursor_description = cursor.description
if conn is not None:
conn.commit()
conn.close()
if query.status == utils.QueryStatus.STOPPED:
return json.dumps(
{
'query_id': query.id,
'status': query.status,
'query': query.to_dict(),
},
default=utils.json_iso_dttm_ser)
cdf = convert_results_to_df(cursor_description, data)
query.rows = cdf.size
query.progress = 100
query.status = QueryStatus.SUCCESS
if query.select_as_cta:
query.select_sql = '{}'.format(
database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema,
show_cols=False,
latest_partition=False))
query.end_time = utils.now_as_float()
session.merge(query)
session.flush()
payload.update({
'status': query.status,
'data': cdf.data if cdf.data else [],
'columns': cdf.columns if cdf.columns else [],
'query': query.to_dict(),
})
if store_results:
key = '{}'.format(uuid.uuid4())
logging.info('Storing results in results backend, key: {}'.format(key))
json_payload = json.dumps(payload, default=utils.json_iso_dttm_ser)
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config.get('CACHE_DEFAULT_TIMEOUT', 0)
results_backend.set(key, utils.zlib_compress(json_payload), cache_timeout)
query.results_key = key
query.end_result_backend_time = utils.now_as_float()
session.merge(query)
session.commit()
if return_results:
return payload
def execute_sql_statements(
ctask, query_id, rendered_query, return_results=True, store_results=False,
user_name=None, session=None, start_time=None,
):
"""Executes the sql query returns the results."""
if store_results and start_time:
# only asynchronous queries
stats_logger.timing(
'sqllab.query.time_pending', now_as_float() - start_time)
query = get_query(query_id, session)
payload = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
if store_results and not results_backend:
raise SqlLabException("Results backend isn't configured.")
# Breaking down into multiple statements
parsed_query = ParsedQuery(rendered_query)
statements = parsed_query.get_statements()
logging.info(f'Executing {len(statements)} statement(s)')
logging.info("Set query to 'running'")
query.status = QueryStatus.RUNNING
query.start_running_time = now_as_float()
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=True,
user_name=user_name,
)
# Sharing a single connection and cursor across the
# execution of all statements (if many)
with closing(engine.raw_connection()) as conn:
with closing(conn.cursor()) as cursor:
statement_count = len(statements)
for i, statement in enumerate(statements):
# TODO CHECK IF STOPPED
msg = f'Running statement {i+1} out of {statement_count}'
logging.info(msg)
query.set_extra_json_key('progress', msg)
session.commit()
is_last_statement = i == len(statements) - 1
try:
cdf = execute_sql_statement(
statement, query, user_name, session, cursor,
return_results=is_last_statement and return_results)
msg = f'Running statement {i+1} out of {statement_count}'
except Exception as e:
msg = str(e)
if statement_count > 1:
msg = f'[Statement {i+1} out of {statement_count}] ' + msg
payload = handle_query_error(msg, query, session, payload)
return payload
# Success, updating the query entry in database
query.rows = cdf.size
query.progress = 100
query.set_extra_json_key('progress', None)
query.status = QueryStatus.SUCCESS
if query.select_as_cta:
query.select_sql = database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema,
show_cols=False,
latest_partition=False)
query.end_time = now_as_float()
session.commit()
payload.update({
'status': query.status,
'data': cdf.data if cdf.data else [],
'columns': cdf.columns if cdf.columns else [],
'query': query.to_dict(),
})
if store_results:
key = str(uuid.uuid4())
logging.info(f'Storing results in results backend, key: {key}')
with stats_timing('sqllab.query.results_backend_write', stats_logger):
json_payload = json.dumps(
payload, default=json_iso_dttm_ser, ignore_nan=True)
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config.get('CACHE_DEFAULT_TIMEOUT', 0)
results_backend.set(key, zlib_compress(json_payload), cache_timeout)
query.results_key = key
session.commit()
if return_results:
return payload
def execute_sql_statements(
ctask,
query_id,
rendered_query,
return_results=True,
store_results=False,
user_name=None,
session=None,
start_time=None,
):
"""Executes the sql query returns the results."""
if store_results and start_time:
# only asynchronous queries
stats_logger.timing('sqllab.query.time_pending',
now_as_float() - start_time)
query = get_query(query_id, session)
payload = dict(query_id=query_id)
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
if store_results and not results_backend:
raise SqlLabException("Results backend isn't configured.")
# Breaking down into multiple statements
parsed_query = ParsedQuery(rendered_query)
statements = parsed_query.get_statements()
logging.info(f'Executing {len(statements)} statement(s)')
logging.info("Set query to 'running'")
query.status = QueryStatus.RUNNING
query.start_running_time = now_as_float()
engine = database.get_sqla_engine(
schema=query.schema,
nullpool=True,
user_name=user_name,
source=sources.get('sql_lab', None),
)
# Sharing a single connection and cursor across the
# execution of all statements (if many)
with closing(engine.raw_connection()) as conn:
with closing(conn.cursor()) as cursor:
statement_count = len(statements)
for i, statement in enumerate(statements):
# TODO CHECK IF STOPPED
msg = f'Running statement {i+1} out of {statement_count}'
logging.info(msg)
query.set_extra_json_key('progress', msg)
session.commit()
try:
cdf = execute_sql_statement(statement, query, user_name,
session, cursor)
msg = f'Running statement {i+1} out of {statement_count}'
except Exception as e:
msg = str(e)
if statement_count > 1:
msg = f'[Statement {i+1} out of {statement_count}] ' + msg
payload = handle_query_error(msg, query, session, payload)
return payload
# Success, updating the query entry in database
query.rows = cdf.size
query.progress = 100
query.set_extra_json_key('progress', None)
query.status = QueryStatus.SUCCESS
if query.select_as_cta:
query.select_sql = database.select_star(
query.tmp_table_name,
limit=query.limit,
schema=database.force_ctas_schema,
show_cols=False,
latest_partition=False)
query.end_time = now_as_float()
payload.update({
'status': query.status,
'data': cdf.data if cdf.data else [],
'columns': cdf.columns if cdf.columns else [],
'query': query.to_dict(),
})
if store_results:
key = str(uuid.uuid4())
logging.info(f'Storing results in results backend, key: {key}')
with stats_timing('sqllab.query.results_backend_write', stats_logger):
json_payload = json.dumps(payload,
default=json_iso_dttm_ser,
ignore_nan=True)
cache_timeout = database.cache_timeout
if cache_timeout is None:
cache_timeout = config.get('CACHE_DEFAULT_TIMEOUT', 0)
results_backend.set(key, zlib_compress(json_payload),
cache_timeout)
query.results_key = key
session.commit()
if return_results:
return payload
def get_sql_results(self, query_id, return_results=True, store_results=False):
"""Executes the sql query returns the results."""
if not self.request.called_directly:
engine = sqlalchemy.create_engine(
app.config.get('SQLALCHEMY_DATABASE_URI'), poolclass=NullPool)
session_class = sessionmaker()
session_class.configure(bind=engine)
session = session_class()
else:
session = db.session()
session.commit() # HACK
query = session.query(models.Query).filter_by(id=query_id).one()
database = query.database
executed_sql = query.sql.strip().strip(';')
db_engine_spec = database.db_engine_spec
def handle_error(msg):
"""Local method handling error while processing the SQL"""
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
raise Exception(query.error_message)
# Limit enforced only for retrieving the data, not for the CTA queries.
is_select = is_query_select(executed_sql);
if not is_select and not database.allow_dml:
handle_error(
"Only `SELECT` statements are allowed against this database")
if query.select_as_cta:
if not is_select:
handle_error(
"Only `SELECT` statements can be used with the CREATE TABLE "
"feature.")
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = 'tmp_{}_table_{}'.format(
query.user_id,
start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
executed_sql = create_table_as(
executed_sql, query.tmp_table_name, database.force_ctas_schema)
query.select_as_cta_used = True
elif (
query.limit and is_select and
db_engine_spec.limit_method == LimitMethod.WRAP_SQL):
executed_sql = database.wrap_sql_limit(executed_sql, query.limit)
query.limit_used = True
engine = database.get_sqla_engine(schema=query.schema)
try:
template_processor = get_template_processor(
database=database, query=query)
executed_sql = template_processor.process_template(executed_sql)
except Exception as e:
logging.exception(e)
msg = "Template rendering failed: " + utils.error_msg_from_exception(e)
handle_error(msg)
try:
query.executed_sql = executed_sql
logging.info("Running query: \n{}".format(executed_sql))
result_proxy = engine.execute(query.executed_sql, schema=query.schema)
except Exception as e:
logging.exception(e)
handle_error(utils.error_msg_from_exception(e))
cursor = result_proxy.cursor
query.status = QueryStatus.RUNNING
session.flush()
db_engine_spec.handle_cursor(cursor, query, session)
cdf = None
if result_proxy.cursor:
column_names = [col[0] for col in result_proxy.cursor.description]
if db_engine_spec.limit_method == LimitMethod.FETCH_MANY:
data = result_proxy.fetchmany(query.limit)
else:
data = result_proxy.fetchall()
cdf = dataframe.SupersetDataFrame(
pd.DataFrame(data, columns=column_names))
query.rows = result_proxy.rowcount
query.progress = 100
query.status = QueryStatus.SUCCESS
if query.rows == -1 and cdf:
# Presto doesn't provide result_proxy.row_count
query.rows = cdf.size
if query.select_as_cta:
query.select_sql = '{}'.format(database.select_star(
query.tmp_table_name, limit=query.limit))
query.end_time = utils.now_as_float()
session.flush()
payload = {
'query_id': query.id,
'status': query.status,
'data': [],
}
payload['data'] = cdf.data if cdf else []
payload['columns'] = cdf.columns_dict if cdf else []
payload['query'] = query.to_dict()
payload = json.dumps(payload, default=utils.json_iso_dttm_ser)
if store_results and results_backend:
key = '{}'.format(uuid.uuid4())
logging.info("Storing results in results backend, key: {}".format(key))
results_backend.set(key, zlib.compress(payload))
query.results_key = key
session.flush()
session.commit()
if return_results:
return payload
