Apache Airflow-使用pymssql + SQLAlchemy

时间:2018-12-04 15:46:50

标签: python sql-server airflow pymssql

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

@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] [驱动程序管理器]找不到数据源名称,并且没有   指定的默认驱动程序(0)(SQLDriverConnect)')(此背景   错误:http://sqlalche.me/e/rvf5

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

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

  

'{conn.conn_type}:// {登录} {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,但无需设置即可:

Airflow UI choose connection type

要解决此问题,我在从 MsSqlHook 继承的新类中重载了 DbApiHook get_uri 方法。我建立了自己的连接字符串,但是根本不干净...

感谢您的帮助

2 个答案:

答案 0 :(得分:2)

您是对的。没有简单,直接的方法让Airflow来完成您想要的事情。我个人将在您的上下文管理器内部构建sqlalchemy引擎,类似于create_engine(hook.get_uri().replace("://", "+pymssql://")) -然后我会将代码扔到可重用的地方。

答案 1 :(得分:0)

您可以通过将连接作为环境变量传递给Airflow来创建连接。参见the docs。变量的值是SqlAlchemy接受的格式的数据库URL。

env var的名称遵循在其后附加连接ID的模式AIRFLOW_CONN_。例如AIRFLOW_CONN_MY_MSSQL,在这种情况下,conn_id将是'my_mssql'