"""
Кастомный 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