来自 Airflow 数据库钩子的 SQLAlchemy 引擎

Question

从 Airflow 连接 ID 获取 SQLAlchemy 引擎的最佳方法是什么？

目前我正在创建一个钩子，检索它的 URI，然后使用它来创建一个 SQLAlchemy 引擎。

postgres_hook = PostgresHook(self.postgres_conn_id)
engine = create_engine(postgres_hook.get_uri())

这有效，但两个命令都连接到数据库。

当我在连接上有“额外”参数时，需要第三个连接来检索这些参数（请参阅从 Airflow Postgres 钩子检索完整连接 URI）

有没有更短更直接的方法？

Answer 1

需要明确的是，您的命令确实会建立两个数据库连接，但它是连接到两个单独的数据库（除非您尝试连接到 Postgres Airflow 数据库）。初始化钩子的第一行不应进行任何连接。只有第二行首先从 Airflow 数据库中获取连接详细信息（我认为您无法避免），然后使用它连接到 Postgres 数据库（我认为这是重点）。

你可以稍微简单一点：

postgres_hook = PostgresHook(self.postgres_conn_id)
engine = postgres_hook.get_sqlalchemy_engine()

这看起来很干净，但如果你想更直接地不经过PostgresHook，你可以通过查询 Airflow 的数据库直接获取它。但是，这意味着您最终将复制代码以从连接对象构建 URI。如果你想继续这个，get_connection()的底层实现是一个很好的例子。

from airflow.settings import Session

conn = session.query(Connection).filter(Connection.conn_id == self.postgres_conn_id).one()
... # build uri from connection
create_engine(uri)

此外，如果您希望能够在extras没有单独的数据库提取的情况下访问超出的内容get_uri()或get_sqlalchemy_engine()功能，您是否可以覆盖BaseHook.get_connection()以将连接对象保存到实例变量以供重用。这需要在 之上创建自己的钩子PostgresHook，所以我知道这可能并不理想。

class CustomPostgresHook(PostgresHook):

    @classmethod
    def get_connection(cls, conn_id):  # type: (str) -> Connection
        conn = super().get_connection(conn_id)
        self.conn_obj = conn  # can't use self.conn because PostgresHook will overriden in https://github.com/apache/airflow/blob/1.10.10/airflow/hooks/postgres_hook.py#L93 by a different type of connection
        return conn

postgres_hook = CustomPostgresHook(self.postgres_conn_id)
uri = postgres_hook.get_uri()
# do something with postgres_hook.conn_obj.extras_dejson

一些内置的 Airflow hooks 已经有这种行为（grpc、samba、tableau），但它绝对不是标准化的。