Загрузка данных


"""
Кастомный Hook для принудительного использования pyodbc
"""
from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook
import pyodbc
from typing import Any


class MsSqlPyodbcHook(MsSqlHook):
    """
    Hook для MS SQL Server с принудительным использованием pyodbc
    """
    
    def get_conn(self) -> pyodbc.Connection:
        """
        Возвращает pyodbc connection напрямую, минуя логику выбора драйвера
        """
        conn = self.get_connection(self.mssql_conn_id)
        
        # Формируем connection string для pyodbc
        server = conn.host or 'localhost'
        database = conn.schema or ''
        username = conn.login or ''
        password = conn.password or ''
        port = conn.port or 1433
        
        # Получаем дополнительные параметры из Extra
        extra_params = conn.extra_dejson
        driver = extra_params.get('driver', 'ODBC Driver 18 for SQL Server')
        trust_cert = extra_params.get('TrustServerCertificate', 'yes')
        
        # Формируем connection string
        conn_str = (
            f"DRIVER={{{driver}}};"
            f"SERVER={server},{port};"
            f"DATABASE={database};"
            f"UID={username};"
            f"PWD={password};"
            f"TrustServerCertificate={trust_cert};"
        )
        
        self.log.info(f'Подключение через pyodbc к {server}:{port}/{database}')
        
        connection = pyodbc.connect(conn_str)
        
        # fast_executemany будет включен на каждом курсоре в _insert_data
        
        return connection
    
    def set_autocommit(self, conn: pyodbc.Connection, autocommit: bool) -> None:
        """
        У pyodbc autocommit - это атрибут, а не метод (в отличие от pymssql)
        """
        conn.autocommit = autocommit
    
    def get_autocommit(self, conn: pyodbc.Connection) -> bool:
        """
        У pyodbc autocommit - это атрибут, у pymssql - autocommit_state
        """
        return conn.autocommit