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

Add close connection, update tests.

Browse source
This commit is contained in:
Pablo Martin 2023-01-20 15:53:45 +01:00
parent 1e0faf702f
commit a15b2f4a86
2 changed files with 15 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

11
lolafect/connections.py
View file

@ -2,6 +2,7 @@ import datetime
import prefect import prefect
from prefect import task from prefect import task
from prefect.triggers import all_finished
from trino.auth import BasicAuthentication from trino.auth import BasicAuthentication
import trino import trino
@ -39,3 +40,13 @@ def connect_to_trino(
) )
return connection 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.")

8
tests/test_integration/test_connections.py
View file

@ -1,7 +1,7 @@
import trino import trino
from lolafect.lolaconfig import build_lolaconfig 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: # This testing suite requires:
# - The calling shell to have permission in AWS # - 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") 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") connection.cursor().execute("SELECT 1")
assert type(connection) == trino.dbapi.Connection close_trino_connection.run(trino_connection=connection)