From 28831bea8c996635bcd425d2f3036e4248958b7e Mon Sep 17 00:00:00 2001 From: Pablo Martin Date: Thu, 19 Jan 2023 17:45:57 +0100 Subject: [PATCH] Improve logging. --- lolafect/connections.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/lolafect/connections.py b/lolafect/connections.py index fc0854a..c2fec0c 100644 --- a/lolafect/connections.py +++ b/lolafect/connections.py @@ -1,5 +1,6 @@ import datetime +import prefect from prefect import task from trino.auth import BasicAuthentication import trino @@ -14,12 +15,15 @@ def connect_to_trino( """ Open a connection to the specified trino instance and return it. - :param trino_credentials: a dict with the host, port, user and password. :param http_schema: which http schema to use in the connection. :return: """ - print("Connecting to Trino.") + logger = prefect.context.get("logger") + logger.info( + f"Connecting to Trino at {trino_credentials['host']}:{trino_credentials['port']}." + ) + connection = trino.dbapi.connect( host=trino_credentials["host"], port=trino_credentials["port"], @@ -30,6 +34,8 @@ def connect_to_trino( trino_credentials["password"], ), ) - print("Connected to Trino.") + logger.info( + f"Connected to Trino at {trino_credentials['host']}:{trino_credentials['port']}." + ) return connection