diff --git a/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py b/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py index 55d59bd57cbfd..6e1ea11c67c70 100644 --- a/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py +++ b/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, diff --git a/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py b/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py new file mode 100644 index 0000000000000..b78d1a6b95763 --- /dev/null +++ b/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")