Question

我正在开发一个自定义气流插件，该插件应显示元数据数据库中的（自定义）表

class CampaignModel(Base, LoggingMixin):
    __tablename__ = "campaigns"

    id = Column(Integer, primary_key=True)
    campaign_name = Column(String(ID_LEN), unique=True)

class CampaignsView(ModelView):
    page_size = 50

v = CampaignsView(CampaignModel, settings.Session(), name="Campaigns")

class CampaignEditorPlugin(AirflowPlugin):
    name = "campaigns_plugin"
    ...
    admin_views = [v]
    ...

Airflow正确启动并加载插件。但是它不会在CampaignModel或airflow initdb上的airflow resetdb中创建/注册我的新表。

如何正确注册表，以便SQLAlchemy在 initdb 上创建表？

Answer 1

尽管在定义类后未通过initdb或resetdb添加CampaignModel.__table__.create(settings.engine, checkfirst=True)会导致在加载模块（如果模块不存在）时创建表。

Answer 2

您可以使用 a) Alembic 迁移或 b) SQLAlchemy ORM 向 Airflow 数据库添加表。

a) 使用 Alembic 进行迁移

向 Airflow 存储库添加代码。

# 1. Create a migration file
cd airflow_dir/airflow/
alembic revision -m "create new_table table"

# 2. Add upgrade and downgrade instructions to the generated migration Python file.
## airflow_dir/airflow/migrations/versions/0a2a5b66e19d_add_new_table_table.py

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = 'xxxxxxxxx'
down_revision = 'xxxxxxxxx'
branch_labels = None
depends_on = None


def upgrade():
    op.create_table('new_table',
        sa.Column('task_id', sa.String(length=250, **COLLATION_ARGS), nullable=False),
        sa.Column('dag_id', sa.String(length=250, **COLLATION_ARGS), nullable=False)

def downgrade():
    op.drop_table('new_table')


# 3. Run the migration
alembic upgrade head

b) SQLAlchemy ORM

可以在与 Airflow 存储库分开的实例目录上的 Airflow 插件中使用。

from airflow import settings
from sqlalchemy import Column, Index, String
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import create_session

class NewModel(Base, LoggingMixin):
    __tablename__ = "new_table"
    id = Column(String(ID_LEN), primary_key=True)
    name = Column(String(ID_LEN), unique=True)


# Create new table in Airflow db
NewModel.__table__.create(settings.engine, checkfirst=True)

# Create object
new_object = NewModel(id='1', name='my name')

# Write to Airflow db
with create_session() as session:
    session.add(new_object)
    session.commit()

Airflow无法从initdb上的自定义数据模型创建表

2 个答案: