feat: refactor all get_sqla_engine to use contextmanager in codebase (#21943)

This commit is contained in:
Hugh A. Miles II
2022-11-15 13:45:14 -05:00
committed by GitHub
parent 06f87e1467
commit e23efefc46
41 changed files with 635 additions and 595 deletions

View File

@@ -162,16 +162,18 @@ class PrestoDBSQLValidator(BaseSQLValidator):
statements = parsed_query.get_statements()
logger.info("Validating %i statement(s)", len(statements))
engine = database.get_sqla_engine(schema, source=QuerySource.SQL_LAB)
# Sharing a single connection and cursor across the
# execution of all statements (if many)
annotations: List[SQLValidationAnnotation] = []
with closing(engine.raw_connection()) as conn:
cursor = conn.cursor()
for statement in parsed_query.get_statements():
annotation = cls.validate_statement(statement, database, cursor)
if annotation:
annotations.append(annotation)
logger.debug("Validation found %i error(s)", len(annotations))
with database.get_sqla_engine_with_context(
schema, source=QuerySource.SQL_LAB
) as engine:
# Sharing a single connection and cursor across the
# execution of all statements (if many)
annotations: List[SQLValidationAnnotation] = []
with closing(engine.raw_connection()) as conn:
cursor = conn.cursor()
for statement in parsed_query.get_statements():
annotation = cls.validate_statement(statement, database, cursor)
if annotation:
annotations.append(annotation)
logger.debug("Validation found %i error(s)", len(annotations))
return annotations