Refactoring and tidy up things.
This commit is contained in:
parent
f4f231d15d
commit
35472c1727
1 changed files with 30 additions and 4 deletions
|
|
@ -62,7 +62,13 @@ def close_trino_connection(trino_connection: trino.dbapi.Connection) -> None:
|
|||
trino_connection.close()
|
||||
logger.info("Trino connection closed successfully.")
|
||||
return
|
||||
logger.info("No connection received.")
|
||||
logger.warning(
|
||||
f"Instead of a Trino connection, a {type(trino_connection)} was received."
|
||||
)
|
||||
raise DeprecationWarning(
|
||||
"This method will only accept the type 'trino.dbapi.Connection' in next major release.\n"
|
||||
"Please, update your code accordingly."
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
|
@ -247,21 +253,41 @@ def connect_to_mysql(
|
|||
|
||||
@task(trigger=all_finished)
|
||||
def close_mysql_connection(connection: pymysql.Connection) -> None:
|
||||
"""
|
||||
Close a MySQL connection, or do nothing if something different is passed.
|
||||
|
||||
:param connection: a MySQL connection.
|
||||
:return: None
|
||||
"""
|
||||
logger = prefect.context.get("logger")
|
||||
|
||||
if isinstance(connection, pymysql.Connection):
|
||||
connection.close()
|
||||
logger.info("DW connection closed successfully.")
|
||||
logger.info("MySQL connection closed successfully.")
|
||||
return
|
||||
logger.info("No connection received.")
|
||||
logger.warning(f"Instead of a MySQL connection, a {type(connection)} was received.")
|
||||
raise DeprecationWarning(
|
||||
"This method will only accept the type 'pymysql.Connection' in next major release.\n"
|
||||
"Please, update your code accordingly."
|
||||
)
|
||||
|
||||
|
||||
@task(trigger=all_finished)
|
||||
def close_ssh_tunnel(tunnel: SSHTunnelForwarder) -> None:
|
||||
"""
|
||||
Close a SSH tunnel, or do nothing if something different is passed.
|
||||
|
||||
:param tunnel: a SSH tunnel.
|
||||
:return:
|
||||
"""
|
||||
logger = prefect.context.get("logger")
|
||||
|
||||
if isinstance(tunnel, SSHTunnelForwarder):
|
||||
tunnel.stop()
|
||||
logger.info("SSH tunnel closed successfully.")
|
||||
return
|
||||
logger.info("No connection received.")
|
||||
logger.warning(f"Instead of a SSH tunnel, a {type(tunnel)} was received.")
|
||||
raise DeprecationWarning(
|
||||
"This method will only accept the type 'SSHTunnelForwarder' in next major release.\n"
|
||||
"Please, update your code accordingly."
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue