Add test_connection method to Trino hook (apache#24583)

airflow/providers/trino/hooks/trino.py

@@ -350,3 +350,19 @@ def insert_rows(
             commit_every = 0
 
         super().insert_rows(table, rows, target_fields, commit_every, replace)
+
+    def test_connection(self):
+        """Tests the connection from UI using Trino specific query"""
+        status, message = False, ''
+        try:
+            with closing(self.get_conn()) as conn:
+                with closing(conn.cursor()) as cur:
+                    cur.execute("select 1")
+                    if cur.fetchone():
+                        status = True
+                        message = 'Connection successfully tested'
+        except Exception as e:
+            status = False
+            message = str(e)
+
+        return status, message

tests/providers/trino/hooks/test_trino.py

@@ -249,6 +249,19 @@ def test_run(self, mock_run):
         self.db_hook.run(sql, autocommit, parameters, list)
         mock_run.assert_called_once_with(sql, autocommit, parameters, handler)
 
+    def test_connection_success(self):
+        status, msg = self.db_hook.test_connection()
+        assert status is True
+        assert msg == 'Connection successfully tested'
+
+    @patch('airflow.providers.trino.hooks.trino.TrinoHook.get_conn')
+    def test_connection_failure(self, mock_conn):
+        mock_conn.side_effect = Exception('Test')
+        self.db_hook.get_conn = mock_conn
+        status, msg = self.db_hook.test_connection()
+        assert status is False
+        assert msg == 'Test'
+
 
 class TestTrinoHookIntegration(unittest.TestCase):
     @pytest.mark.integration("trino")