From dc9895f8540d7e5335419c047097c5b3fcc85784 Mon Sep 17 00:00:00 2001 From: Shubham Kapoor Date: Mon, 24 Aug 2026 19:55:36 +0530 Subject: [PATCH] IBM Db2 provider: skip None-valued extra parameters to avoid KEY=None in connection strings When a user leaves an optional extra field blank in the Airflow connection form, the JSON stored in the extra field contains null values (e.g. {"SSLServerCertificate": null}). Previously these were emitted verbatim into both the ibm_db connection string (KEY=None;) and the SQLAlchemy URI query string (?KEY=None), causing the Db2 driver to receive the literal string "None" as a parameter value. This either triggers a connection error or silently passes a bad value to the driver. Fix: skip any extra key whose value is None before building the connection string in get_conn() and before building the query string in get_uri(). Add parametrized tests covering both methods to prevent regression. --- .../airflow/providers/ibm/db2/hooks/db2.py | 4 ++ .../db2/tests/unit/ibm/db2/hooks/test_db2.py | 60 +++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/providers/ibm/db2/src/airflow/providers/ibm/db2/hooks/db2.py b/providers/ibm/db2/src/airflow/providers/ibm/db2/hooks/db2.py index 314b506285e3e..6db8f176764ae 100644 --- a/providers/ibm/db2/src/airflow/providers/ibm/db2/hooks/db2.py +++ b/providers/ibm/db2/src/airflow/providers/ibm/db2/hooks/db2.py @@ -99,6 +99,8 @@ def get_conn(self) -> Any: # Add all extra parameters to connection string # Parameter names are automatically converted to uppercase for Db2 for key, value in extra.items(): + if value is None: + continue # Convert boolean values to appropriate strings if isinstance(value, bool): converted_value = "true" if value else "false" @@ -134,6 +136,8 @@ def get_uri(self) -> str: if extra: query_params = {} for key, value in extra.items(): + if value is None: + continue # Convert boolean values to appropriate strings if isinstance(value, bool): query_params[key.upper()] = "true" if value else "false" diff --git a/providers/ibm/db2/tests/unit/ibm/db2/hooks/test_db2.py b/providers/ibm/db2/tests/unit/ibm/db2/hooks/test_db2.py index 5a179bb0fb2d4..817f9e5e5380b 100644 --- a/providers/ibm/db2/tests/unit/ibm/db2/hooks/test_db2.py +++ b/providers/ibm/db2/tests/unit/ibm/db2/hooks/test_db2.py @@ -99,6 +99,38 @@ def test_get_conn_with_ssl(self, mock_get_connection, mock_connection_with_extra assert "SECURITY=SSL" in call_args assert "SSLSERVERCERTIFICATE=/path/to/cert.crt" in call_args + @pytest.mark.parametrize( + ("extra", "expected_absent"), + [ + ('{"SSLServerCertificate": null}', ["SSLSERVERCERTIFICATE=None", "SSLSERVERCERTIFICATE="]), + ('{"SECURITY": "SSL", "SSLServerCertificate": null}', ["SSLSERVERCERTIFICATE=None"]), + ], + ) + @patch("airflow.providers.ibm.db2.hooks.db2.Db2Hook.get_connection") + def test_get_conn_skips_none_extra_values(self, mock_get_connection, extra, expected_absent): + """None-valued extra keys must not be emitted in the Db2 connection string.""" + conn = Connection( + conn_id="db2_default", + conn_type="db2", + host="localhost", + login="db2user", + password="db2pass", + schema="testdb", + port=50000, + extra=extra, + ) + mock_get_connection.return_value = conn + mock_ibm_db_dbi = MagicMock() + mock_ibm_db_dbi.connect.return_value = MagicMock() + + with patch.dict(sys.modules, {"ibm_db_dbi": mock_ibm_db_dbi}): + hook = Db2Hook(db2_conn_id="db2_default") + hook.get_conn() + + call_args = mock_ibm_db_dbi.connect.call_args[0][0] + for absent in expected_absent: + assert absent not in call_args + @patch("airflow.providers.ibm.db2.hooks.db2.Db2Hook.get_connection") def test_get_uri(self, mock_get_connection, mock_connection): """Test get_uri method.""" @@ -150,6 +182,34 @@ def test_get_uri_with_defaults(self, mock_get_connection): # Verify URI uses defaults assert uri == "db2+ibm_db://:@localhost:50000/" + @pytest.mark.parametrize( + ("extra", "expected_absent"), + [ + ('{"SSLServerCertificate": null}', ["SSLSERVERCERTIFICATE=None", "SSLSERVERCERTIFICATE="]), + ('{"SECURITY": "SSL", "SSLServerCertificate": null}', ["SSLSERVERCERTIFICATE=None"]), + ], + ) + @patch("airflow.providers.ibm.db2.hooks.db2.Db2Hook.get_connection") + def test_get_uri_skips_none_extra_values(self, mock_get_connection, extra, expected_absent): + """None-valued extra keys must not be emitted in the SQLAlchemy URI query string.""" + conn = Connection( + conn_id="db2_default", + conn_type="db2", + host="localhost", + login="db2user", + password="db2pass", + schema="testdb", + port=50000, + extra=extra, + ) + mock_get_connection.return_value = conn + + hook = Db2Hook(db2_conn_id="db2_default") + uri = hook.get_uri() + + for absent in expected_absent: + assert absent not in uri + def test_hook_attributes(self): """Test hook class attributes.""" assert Db2Hook.conn_name_attr == "db2_conn_id"