diff --git a/lolafect/connections.py b/lolafect/connections.py index 450c4e4..caf3fdf 100644 --- a/lolafect/connections.py +++ b/lolafect/connections.py @@ -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." + )