From a15b2f4a861fbbd5c0a79ed4c9a9af353cc62c52 Mon Sep 17 00:00:00 2001 From: Pablo Martin Date: Fri, 20 Jan 2023 15:53:45 +0100 Subject: [PATCH] Add close connection, update tests. --- lolafect/connections.py | 11 +++++++++++ tests/test_integration/test_connections.py | 8 ++++---- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/lolafect/connections.py b/lolafect/connections.py index c2fec0c..eb14237 100644 --- a/lolafect/connections.py +++ b/lolafect/connections.py @@ -2,6 +2,7 @@ import datetime import prefect from prefect import task +from prefect.triggers import all_finished from trino.auth import BasicAuthentication import trino @@ -39,3 +40,13 @@ def connect_to_trino( ) return connection + + +@task(trigger=all_finished) +def close_trino_connection(trino_connection: trino.dbapi.Connection): + logger = prefect.context.get("logger") + if isinstance(trino_connection, trino.dbapi.Connection): + trino_connection.close() + logger.info("Trino connection closed successfully.") + return + logger.info("No connection received.") diff --git a/tests/test_integration/test_connections.py b/tests/test_integration/test_connections.py index 5b79a46..bbf8ddd 100644 --- a/tests/test_integration/test_connections.py +++ b/tests/test_integration/test_connections.py @@ -1,7 +1,7 @@ import trino from lolafect.lolaconfig import build_lolaconfig -from lolafect.connections import connect_to_trino +from lolafect.connections import connect_to_trino, close_trino_connection # This testing suite requires: # - The calling shell to have permission in AWS @@ -9,10 +9,10 @@ from lolafect.connections import connect_to_trino TEST_LOLACONFIG = build_lolaconfig(flow_name="testing-suite") -def test_that_trino_connector_works_properly(): +def test_that_trino_connect_and_disconnect_works_properly(): - connection = connect_to_trino.run(TEST_LOLACONFIG.TRINO_CREDENTIALS) + connection = connect_to_trino.run(trino_credentials=TEST_LOLACONFIG.TRINO_CREDENTIALS) connection.cursor().execute("SELECT 1") - assert type(connection) == trino.dbapi.Connection + close_trino_connection.run(trino_connection=connection) \ No newline at end of file