counterweight/lolafect
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Small details

Browse source
This commit is contained in:
Pablo Martin 2023-01-23 18:38:52 +01:00
parent 8c66a843b9
commit e2e7f8fb93
1 changed files with 6 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

10
lolafect/connections.py
View file

@ -215,7 +215,7 @@ def connect_to_mysql(
mysql_credentials["port"], mysql_credentials["port"],
), ),
local_bind_address=("127.0.0.1", mysql_credentials["port"]), local_bind_address=("127.0.0.1", mysql_credentials["port"]),
ssh_private_key_password=mysql_credentials["ssh_pkey_password"], ssh_private_key_password=ssh_tunnel_credentials["ssh_pkey_password"],
) )
except Exception as e: except Exception as e:
raise e raise e
@ -241,13 +241,15 @@ def connect_to_mysql(
db_connection.raw_user = mysql_credentials["user"] db_connection.raw_user = mysql_credentials["user"]
db_connection.raw_password = mysql_credentials["password"] db_connection.raw_password = mysql_credentials["password"]
logger.info(f"Connected to MySQL at {mysql_credentials['host']}:{mysql_credentials['port']}.") logger.info(
f"Connected to MySQL at {mysql_credentials['host']}:{mysql_credentials['port']}."
)
return db_connection, tunnel return db_connection, tunnel
@task(trigger=all_finished) @task(trigger=all_finished)
def close_dw_connection(dw_connection: pymysql.Connection) -> None: def close_mysql_connection(dw_connection: pymysql.Connection) -> None:
logger = prefect.context.get("logger") logger = prefect.context.get("logger")
if isinstance(dw_connection, pymysql.Connection): if isinstance(dw_connection, pymysql.Connection):
@ -257,7 +259,7 @@ def close_dw_connection(dw_connection: pymysql.Connection) -> None:
logger.info("No connection received.") logger.info("No connection received.")
@task(log_stdout=True, trigger=all_finished) @task(trigger=all_finished)
def close_ssh_tunnel(ssh_tunnel: SSHTunnelForwarder) -> None: def close_ssh_tunnel(ssh_tunnel: SSHTunnelForwarder) -> None:
logger = prefect.context.get("logger") logger = prefect.context.get("logger")