Skip to content

Commit

Permalink
fix: don't close cursor before closing connection
Browse files Browse the repository at this point in the history
  • Loading branch information
betodealmeida committed Jan 28, 2021
1 parent 365770e commit eb3fb37
Showing 1 changed file with 35 additions and 37 deletions.
72 changes: 35 additions & 37 deletions superset/sql_lab.py
Original file line number Diff line number Diff line change
Expand Up @@ -379,44 +379,42 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca
# 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

# 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
)
)
# 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

# 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)
if statement_count > 1:
msg = f"[Statement {i+1} out of {statement_count}] " + msg
payload = handle_query_error(msg, query, session, payload)
return payload
# 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)
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()
Expand Down

0 comments on commit eb3fb37

Please sign in to comment.