Add trino task.

lolafect/connections.py

@@ -0,0 +1,35 @@
+import datetime
+
+from prefect import task
+from trino.auth import BasicAuthentication
+import trino
+
+from lolafect.defaults import DEFAULT_TRINO_HTTP_SCHEME
+
+
+@task(log_stdout=True, max_retries=3, retry_delay=datetime.timedelta(minutes=10))
+def connect_to_trino(
+    trino_credentials: dict, http_schema: str = DEFAULT_TRINO_HTTP_SCHEME
+) -> trino.dbapi.Connection:
+    """
+    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.")
+    connection = trino.dbapi.connect(
+        host=trino_credentials["host"],
+        port=trino_credentials["port"],
+        user=trino_credentials["user"],
+        http_scheme=http_schema,
+        auth=BasicAuthentication(
+            trino_credentials["user"],
+            trino_credentials["password"],
+        ),
+    )
+    print("Connected to Trino.")
+
+    return connection

lolafect/defaults.py

@@ -4,3 +4,4 @@ DEFAULT_PATH_TO_SLACK_WEBHOOKS_FILE = "env/slack_webhooks.json"
 DEFAULT_KUBERNETES_IMAGE = "373245262072.dkr.ecr.eu-central-1.amazonaws.com/pdo-data-prefect:production"
 DEFAULT_KUBERNETES_LABELS  = ["k8s"]
 DEFAULT_FLOWS_PATH_IN_BUCKET = "flows/"
+DEFAULT_TRINO_HTTP_SCHEME = "https"