python - Apache Airflow - 使用 pymssql + SQLAlchemy 与 MS SQL Server 的连接问题

Question

我在使用 pymssql 连接到 Apache Airflow 1.10.1 中的 Azure MS SQL Server 2014 数据库时遇到问题。我想使用 Airflow 提供的 MsSqlHook 类，以便于在 Airflow UI 中创建连接，然后使用 SqlAlchemy 为我的连接创建上下文管理器:

@contextmanager
def mssql_session(dt_conn_id):
    sqla_engine = MsSqlHook(mssql_conn_id=dt_conn_id).get_sqlalchemy_engine()
    session = sessionmaker(bind=sqla_engine)()
    try:
        yield session
    except:
        session.rollback()
        raise
    else:
        session.commit()
    finally:
        session.close()

但是当我这样做时，我在运行请求时遇到此错误:

sqlalchemy.exc.InterfaceError: (pyodbc.InterfaceError) ('IM002', '[IM002] [unixODBC][Driver Manager]Data source name not found, and no default driver specified (0) (SQLDriverConnect)') (Background on this error at: http://sqlalche.me/e/rvf5)

它似乎来自pyodbc，而我想使用pymssql(在MsSqlHook中，方法get_conn使用pymssql!)

我在Airflow的源代码中查找了原因。我注意到类 DbApiHook 中的方法 get_uri(从中继承 MsSqlHook)构建了传递给 SqlAlchemy 的连接字符串em>像这样:

'{conn.conn_type}://{login}{host}/{conn.schema}'

但是 conn.conn_type 只是等于“mssql”，而我们需要指定 DBAPI，如下所述: https://docs.sqlalchemy.org/en/latest/core/engines.html#microsoft-sql-server(例如:'mssql+pymssql://scott:tiger@hostname:port/dbname')

所以，默认情况下，我认为它使用pyodbc。但是如何正确设置连接的 conn_type 为 'mssql+pymssql' 而不是 'mssql' ？在Airflow IU中，您可以简单地在下拉列表中选择SQL server，但不能根据需要进行设置:

为了解决此问题，我在从 MsSqlHook 创建的继承自 MsSqlHook 的新类中重载了 DbApiHook 中的 get_uri 方法，其中我构建了自己的连接字符串，但它根本不干净......

感谢您的帮助

Answer 1

你是对的。没有简单、直接的方法可以让 Airflow 执行您想要的操作。就我个人而言，我会在上下文管理器中构建 sqlalchemy 引擎，例如 create_engine(hook.get_uri().replace("://", "+pymssql://")) ——然后我会把代码扔到可重用的地方。