From 4dd94572d1c1b6bdbe5131c3f830aa66b3e867ec Mon Sep 17 00:00:00 2001 From: Piyush Date: Sun, 21 Sep 2025 04:15:52 +0000 Subject: [PATCH 1/6] fix: vertica sqlalchemy uri --- airflow-core/newsfragments/38195.bugfix.rst | 1 + .../providers/vertica/hooks/vertica.py | 29 +- .../unit/vertica/hooks/test_vertica_sql.py | 272 ++++++++++++++++++ 3 files changed, 301 insertions(+), 1 deletion(-) create mode 100644 airflow-core/newsfragments/38195.bugfix.rst create mode 100644 providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py diff --git a/airflow-core/newsfragments/38195.bugfix.rst b/airflow-core/newsfragments/38195.bugfix.rst new file mode 100644 index 0000000000000..bb12513f3096e --- /dev/null +++ b/airflow-core/newsfragments/38195.bugfix.rst @@ -0,0 +1 @@ +Fix a bug where the VerticaHook could not properly parse SQLAlchemy URIs. \ No newline at end of file diff --git a/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py b/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py index 55d59bd57cbfd..e985e8f9b98d6 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): + """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): + """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..a601de81da4e9 --- /dev/null +++ b/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py @@ -0,0 +1,272 @@ +# 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.mock import MagicMock, PropertyMock, call, 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, + ) + + actual_calls = [call.args[0] for call in mock_cursor.execute.call_args_list] + assert actual_calls == expected_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( + [ + call("SELECT 1;"), + call("SELECT 2;"), + ] + ) + + assert results == [(2,)] + + +def test_get_uri(vertica_hook): + """ + Test that the get_uri() method returns the correct connection string. + """ + + expected_uri = "vertica-python://user:pass@vertica.cloud.com:5433/test_db" + + actual_uri = vertica_hook.get_uri() + + assert actual_uri == expected_uri + + +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") From dc38cd23f2ad8da88e49625d5aaf00af17ec11fb Mon Sep 17 00:00:00 2001 From: Piyush Date: Sun, 21 Sep 2025 07:23:37 +0000 Subject: [PATCH 2/6] Fix: add newline at end of 38195.bugfix.rst to pass pre-commit hooks --- airflow-core/newsfragments/38195.bugfix.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/newsfragments/38195.bugfix.rst b/airflow-core/newsfragments/38195.bugfix.rst index bb12513f3096e..c35f087bd4f8a 100644 --- a/airflow-core/newsfragments/38195.bugfix.rst +++ b/airflow-core/newsfragments/38195.bugfix.rst @@ -1 +1 @@ -Fix a bug where the VerticaHook could not properly parse SQLAlchemy URIs. \ No newline at end of file +Fix a bug where the VerticaHook could not properly parse SQLAlchemy URIs. From 8ef8897e21661773ae803f6cd3e6407230386bc3 Mon Sep 17 00:00:00 2001 From: Piyush Date: Sun, 21 Sep 2025 07:31:46 +0000 Subject: [PATCH 3/6] Fix: Deleted the bugfix --- airflow-core/newsfragments/38195.bugfix.rst | 1 - 1 file changed, 1 deletion(-) delete mode 100644 airflow-core/newsfragments/38195.bugfix.rst diff --git a/airflow-core/newsfragments/38195.bugfix.rst b/airflow-core/newsfragments/38195.bugfix.rst deleted file mode 100644 index c35f087bd4f8a..0000000000000 --- a/airflow-core/newsfragments/38195.bugfix.rst +++ /dev/null @@ -1 +0,0 @@ -Fix a bug where the VerticaHook could not properly parse SQLAlchemy URIs. From d3d9eb0c62a3ff5048b10dc8a43477c4f27cbb7b Mon Sep 17 00:00:00 2001 From: Piyush Date: Sun, 21 Sep 2025 11:58:37 +0000 Subject: [PATCH 4/6] Added the annotations --- .../vertica/src/airflow/providers/vertica/hooks/vertica.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py b/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py index e985e8f9b98d6..6e1ea11c67c70 100644 --- a/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py +++ b/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py @@ -135,7 +135,7 @@ def get_conn(self) -> connect: return conn @property - def sqlalchemy_url(self): + 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 {} @@ -157,7 +157,7 @@ def sqlalchemy_url(self): query=query, ) - def get_uri(self): + def get_uri(self) -> str: """Return a URI string with password visible.""" return self.sqlalchemy_url.render_as_string(hide_password=False) From 135e12cfcc1edc027436cf3a0e555b92832ce6c8 Mon Sep 17 00:00:00 2001 From: Piyush Date: Mon, 22 Sep 2025 14:09:14 +0000 Subject: [PATCH 5/6] adjusting the format for mocks --- .../vertica/tests/unit/vertica/hooks/test_vertica_sql.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py b/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py index a601de81da4e9..b202fd7626f42 100644 --- a/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py +++ b/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py @@ -160,8 +160,8 @@ def test_vertica_run_queries( return_last=return_last, ) - actual_calls = [call.args[0] for call in mock_cursor.execute.call_args_list] - assert actual_calls == expected_calls + expected_mock_calls = [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 @@ -194,6 +194,7 @@ def test_get_uri(vertica_hook): actual_uri = vertica_hook.get_uri() assert actual_uri == expected_uri + assert vertica_hook.get_uri() == "vertica-python://user:pass@vertica.cloud.com:5433/test_db" def test_get_sqlalchemy_engine(vertica_hook): From 3a9d8128212235d5870aa27d1075da389771d38f Mon Sep 17 00:00:00 2001 From: Piyush Date: Fri, 26 Sep 2025 08:03:57 +0000 Subject: [PATCH 6/6] Changed the Mock Format --- .../tests/unit/vertica/hooks/test_vertica_sql.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py b/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py index b202fd7626f42..b78d1a6b95763 100644 --- a/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py +++ b/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py @@ -19,7 +19,8 @@ import json from collections import namedtuple -from unittest.mock import MagicMock, PropertyMock, call, patch +from unittest import mock +from unittest.mock import MagicMock, PropertyMock, patch import pytest from sqlalchemy.engine import Engine @@ -160,7 +161,7 @@ def test_vertica_run_queries( return_last=return_last, ) - expected_mock_calls = [call(sql_call) for sql_call in expected_calls] + 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 @@ -176,8 +177,8 @@ def test_run_with_multiple_statements(vertica_hook, mock_cursor): mock_cursor.execute.assert_has_calls( [ - call("SELECT 1;"), - call("SELECT 2;"), + mock.call("SELECT 1;"), + mock.call("SELECT 2;"), ] ) @@ -188,12 +189,6 @@ def test_get_uri(vertica_hook): """ Test that the get_uri() method returns the correct connection string. """ - - expected_uri = "vertica-python://user:pass@vertica.cloud.com:5433/test_db" - - actual_uri = vertica_hook.get_uri() - - assert actual_uri == expected_uri assert vertica_hook.get_uri() == "vertica-python://user:pass@vertica.cloud.com:5433/test_db"