Amitkumar293 commented on code in PR #67532: URL: https://github.com/apache/airflow/pull/67532#discussion_r3396161587
########## providers/ibm/db2/src/airflow/providers/ibm/db2/hooks/db2.py: ########## @@ -0,0 +1,186 @@ +# +# 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. +"""This module allows to connect to an IBM Db2 database.""" + +from __future__ import annotations + +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.sql.hooks.sql import DbApiHook +from airflow.providers.ibm.db2.dialects.db2 import Db2Dialect + +if TYPE_CHECKING: + from airflow.providers.common.sql.dialects.dialect import Dialect + + +class Db2Hook(DbApiHook): + """ + Interact with IBM Db2. + + You can specify SSL and other connection parameters in the extra field of your connection. + + **SSL Configuration:** + - ``{"SECURITY": "SSL", "SSLServerCertificate": "/path/to/cert.crt"}`` - For SSL with server certificate + - ``{"SECURITY": "SSL"}`` - For SSL without certificate validation + + **Other Parameters:** + Any additional Db2 connection string parameters can be added to the extra field. + Parameter names will be automatically converted to uppercase. + + :param db2_conn_id: The :ref:`Db2 connection id <howto/connection:Db2>` + reference to a specific Db2 database. + """ + + conn_name_attr = "db2_conn_id" + default_conn_name = "db2_default" + conn_type = "Db2" + hook_name = "IBM Db2" + supports_autocommit = True + supports_executemany = True + _test_connection_sql = "SELECT 1 FROM SYSIBM.SYSDUMMY1" + _placeholder = "?" + + @cached_property + def dialect_name(self) -> str: + return "db2" + + @cached_property + def dialect(self) -> Dialect: + return Db2Dialect(self) + + @staticmethod + def _get_default_values(conn) -> dict[str, Any]: + """Extract connection values with defaults.""" + return { + "host": conn.host or "localhost", + "port": conn.port or 50000, + "schema": conn.schema or "", + "login": conn.login or "", + "password": conn.password or "", + } + + def get_conn(self) -> Any: + """ + Return ibm_db_dbi connection object. + + :return: Db2 connection object + """ + import ibm_db_dbi + + conn = self.get_connection(self.get_conn_id()) + defaults = self._get_default_values(conn) + + # Build connection string for Db2 + conn_str_parts = [ + f"DATABASE={defaults['schema']}", + f"HOSTNAME={defaults['host']}", + f"PORT={defaults['port']}", + "PROTOCOL=TCPIP", + f"UID={defaults['login']}", + f"PWD={defaults['password']}", + ] + + # Add extra connection parameters from connection extra field + extra = conn.extra_dejson + if extra: + # Add all extra parameters to connection string + # Parameter names are automatically converted to uppercase for Db2 + for key, value in extra.items(): + # Convert boolean values to appropriate strings + if isinstance(value, bool): + converted_value = "true" if value else "false" + else: + converted_value = value + conn_str_parts.append(f"{key.upper()}={converted_value}") + + conn_str = ";".join(conn_str_parts) + ";" Review Comment: Thank you for the feedback! I've implemented the suggested refactoring. The connection string building logic is now extracted into a dedicated get_conn_str() method, making get_conn() much cleaner: -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
