Fix: Handle SQLAlchemy_URI for Vertica Provider

providers/vertica/src/airflow/providers/vertica/hooks/vertica.py

@@ -1,4 +1,3 @@
-#
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
@@ -20,6 +19,7 @@
 from collections.abc import Callable, Iterable, Mapping
 from typing import Any, overload
 
+from sqlalchemy.engine import URL
 from vertica_python import connect
 
 from airflow.providers.common.sql.hooks.handlers import fetch_all_handler
@@ -134,6 +134,33 @@ def get_conn(self) -> connect:
         conn = connect(**conn_config)
         return conn
 
+    @property
+    def sqlalchemy_url(self) -> URL:
+        """Return a SQLAlchemy URL object with properly formatted query parameters."""
+        conn = self.get_connection(self.get_conn_id())
+        extra = conn.extra_dejson or {}
+
+        # Normalize query dictionary
+        query = {
+            k: ([str(x) for x in v] if isinstance(v, (list, tuple)) else str(v))
+            for k, v in extra.items()
+            if v is not None
+        }
+
+        return URL.create(
+            drivername="vertica-python",
+            username=conn.login,
+            password=conn.password or "",
+            host=conn.host or "localhost",
+            port=conn.port or 5433,
+            database=conn.schema,
+            query=query,
+        )
+
+    def get_uri(self) -> str:
+        """Return a URI string with password visible."""
+        return self.sqlalchemy_url.render_as_string(hide_password=False)
+
     @overload
     def run(
         self,

providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py

@@ -0,0 +1,268 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import json
+from collections import namedtuple
+from unittest import mock
+from unittest.mock import MagicMock, PropertyMock, patch
+
+import pytest
+from sqlalchemy.engine import Engine
+
+from airflow.exceptions import AirflowException
+from airflow.models import Connection
+from airflow.providers.vertica.hooks.vertica import VerticaHook
+
+DEFAULT_CONN_ID = "vertica_default"
+HOST = "vertica.cloud.com"
+PORT = 5433
+USER = "user"
+PASSWORD = "pass"
+DATABASE = "test_db"
+
+SerializableRow = namedtuple("SerializableRow", ["id", "value"])
+
+
+def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]:
+    """Convert field names into cursor.description tuples."""
+    return [(field,) for field in fields]
+
+
+@pytest.fixture(autouse=True)
+def create_connection(create_connection_without_db):
+    """Create a mocked Airflow connection for Vertica."""
+    create_connection_without_db(
+        Connection(
+            conn_id=DEFAULT_CONN_ID,
+            conn_type="vertica",
+            host=HOST,
+            login=USER,
+            password=PASSWORD,
+            schema=DATABASE,
+        )
+    )
+
+
+@pytest.fixture
+def vertica_hook():
+    return VerticaHook(vertica_conn_id=DEFAULT_CONN_ID)
+
+
+@pytest.fixture
+def mock_get_conn():
+    with patch("airflow.providers.vertica.hooks.vertica.VerticaHook.get_conn") as mock_conn:
+        yield mock_conn
+
+
+@pytest.fixture
+def mock_cursor(mock_get_conn):
+    cursor = MagicMock()
+    type(cursor).rowcount = PropertyMock(return_value=1)
+    cursor.fetchall.return_value = [("1", "row1")]
+    cursor.description = get_cursor_descriptions(["id", "value"])
+    cursor.nextset.side_effect = [False]
+    mock_get_conn.return_value.cursor.return_value = cursor
+    return cursor
+
+
+def test_sqlalchemy_url_property(vertica_hook):
+    url = vertica_hook.sqlalchemy_url.render_as_string(hide_password=False)
+    expected_url = f"vertica-python://{USER}:{PASSWORD}@{HOST}:{PORT}/{DATABASE}"
+    assert url.startswith(expected_url)
+
+
+@pytest.mark.parametrize(
+    "return_last, split_statements, sql, expected_calls, cursor_results, expected_result",
+    [
+        pytest.param(
+            True,
+            False,
+            "SELECT * FROM table",
+            ["SELECT * FROM table"],
+            [("1", "row1"), ("2", "row2")],
+            [SerializableRow("1", "row1"), SerializableRow("2", "row2")],
+            id="Single query, return_last=True",
+        ),
+        pytest.param(
+            False,
+            False,
+            "SELECT * FROM table",
+            ["SELECT * FROM table"],
+            [("1", "row1"), ("2", "row2")],
+            [SerializableRow("1", "row1"), SerializableRow("2", "row2")],
+            id="Single query, return_last=False",
+        ),
+        pytest.param(
+            True,
+            True,
+            "SELECT * FROM table1; SELECT * FROM table2;",
+            ["SELECT * FROM table1;", "SELECT * FROM table2;"],
+            [[("1", "row1"), ("2", "row2")], [("3", "row3"), ("4", "row4")]],
+            [SerializableRow("3", "row3"), SerializableRow("4", "row4")],
+            id="Multiple queries, split_statements=True, return_last=True",
+        ),
+        pytest.param(
+            True,
+            False,
+            "SELECT * FROM table1; SELECT * FROM table2;",
+            ["SELECT * FROM table1; SELECT * FROM table2;"],
+            [("1", "row1"), ("2", "row2")],
+            [SerializableRow("1", "row1"), SerializableRow("2", "row2")],
+            id="Multiple queries, split_statements=False",
+        ),
+        pytest.param(
+            True,
+            False,
+            "SELECT * FROM empty",
+            ["SELECT * FROM empty"],
+            [],
+            [],
+            id="Empty result",
+        ),
+    ],
+)
+def test_vertica_run_queries(
+    vertica_hook,
+    mock_cursor,
+    return_last,
+    split_statements,
+    sql,
+    expected_calls,
+    cursor_results,
+    expected_result,
+):
+    if split_statements:
+        mock_cursor.fetchall.side_effect = cursor_results
+        mock_cursor.nextset.side_effect = [True] * (len(cursor_results) - 1) + [False]
+    else:
+        mock_cursor.fetchall.return_value = cursor_results
+        mock_cursor.nextset.side_effect = lambda: False
+
+    result = vertica_hook.run(
+        sql,
+        handler=lambda cur: cur.fetchall(),
+        split_statements=split_statements,
+        return_last=return_last,
+    )
+
+    expected_mock_calls = [mock.call(sql_call) for sql_call in expected_calls]
+    mock_cursor.execute.assert_has_calls(expected_mock_calls)
+    assert [SerializableRow(*row) for row in result] == expected_result
+
+
+def test_run_with_multiple_statements(vertica_hook, mock_cursor):
+    mock_cursor.fetchall.side_effect = [[(1,)], [(2,)]]
+
+    mock_cursor.nextset.side_effect = [True, False]
+
+    sql = "SELECT 1; SELECT 2;"
+
+    results = vertica_hook.run(sql, handler=lambda cur: cur.fetchall(), split_statements=True)
+
+    mock_cursor.execute.assert_has_calls(
+        [
+            mock.call("SELECT 1;"),
+            mock.call("SELECT 2;"),
+        ]
+    )
+
+    assert results == [(2,)]
+
+
+def test_get_uri(vertica_hook):
+    """
+    Test that the get_uri() method returns the correct connection string.
+    """
+    assert vertica_hook.get_uri() == "vertica-python://user:pass@vertica.cloud.com:5433/test_db"
+
+
+def test_get_sqlalchemy_engine(vertica_hook):
+    """
+    Test that the get_sqlalchemy_engine() method returns a valid SQLAlchemy engine.
+    """
+
+    with patch("airflow.providers.common.sql.hooks.sql.create_engine") as mock_create_engine:
+        mock_engine = MagicMock(spec=Engine)
+        mock_create_engine.return_value = mock_engine
+
+        engine = vertica_hook.get_sqlalchemy_engine()
+
+        assert engine is mock_engine
+
+        mock_create_engine.assert_called_once()
+        call_args = mock_create_engine.call_args[1]
+        assert "url" in call_args
+
+        actual_url = call_args["url"]
+
+        assert actual_url.drivername == "vertica-python"
+        assert actual_url.username == "user"
+        assert actual_url.password == "pass"
+        assert actual_url.host == "vertica.cloud.com"
+        assert actual_url.port == 5433
+        assert actual_url.database == "test_db"
+
+
+@pytest.mark.parametrize("sql", ["", "\n", " "])
+def test_run_with_no_query(vertica_hook, sql):
+    """
+    Test that running with no SQL query raises a ValueError.
+    """
+    with pytest.raises(ValueError, match="List of SQL statements is empty"):
+        vertica_hook.run(sql)
+
+
+def test_run_with_invalid_column_names(vertica_hook, mock_cursor):
+    invalid_names = [("1_2_3",), ("select",), ("from",)]
+    mock_cursor.description = invalid_names
+    mock_cursor.fetchall.return_value = [(1, "row1", "bar")]
+
+    result = vertica_hook.run(sql="SELECT * FROM table", handler=lambda cur: cur.fetchall())
+
+    assert result[0][0] == 1
+    assert result[0][1] == "row1"
+    assert result[0][2] == "bar"
+
+
+@pytest.fixture
+def vertica_hook_with_timeout(create_connection_without_db):
+    create_connection_without_db(
+        Connection(
+            conn_id="vertica_timeout",
+            conn_type="vertica",
+            host="vertica.cloud.com",
+            login="user",
+            password="pass",
+            schema="test_db",
+            extra=json.dumps({"execution_timeout": 1}),
+        )
+    )
+    return VerticaHook(vertica_conn_id="vertica_timeout")
+
+
+def test_execution_timeout_exceeded(vertica_hook_with_timeout, mock_cursor):
+    mock_error_response = MagicMock()
+    mock_error_response.error_message.return_value = "Mock error message for test."
+
+    with patch(
+        "airflow.providers.common.sql.hooks.sql.DbApiHook.run",
+        side_effect=AirflowException("Query exceeded execution timeout"),
+    ):
+        with pytest.raises(AirflowException, match="Query exceeded execution timeout"):
+            vertica_hook_with_timeout.run(sql="SELECT * FROM table1")