Repository: incubator-airflow Updated Branches: refs/heads/v1-10-test 80fd31250 -> 54da3eaf1
[AIRFLOW-2606] Fix DB schema and SQLAlchemy model * Add test that verifies that database schema and SQLAlchemy model are in sync * Add exception for users.password that doesn't exist in model and tables created by other tests * Add migration script to merge the two heads * Add migration script to fix not-null constrains for MySQL that were lost by 0e2a74e0fc9f_add_time_zone_awareness * Add migration script to fix FK constraint for existing SQLite DBs * Enable ForeignKey support for SQLite, otherwise 2e82aab8ef20_rename_user_table won't change FK in chart and known_event tables (cherry picked from commit 680651f0ae2a314f8e9882a6bc38f4fa3795cdbe) Signed-off-by: Bolke de Bruin <[email protected]> Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/54da3eaf Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/54da3eaf Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/54da3eaf Branch: refs/heads/v1-10-test Commit: 54da3eaf175b0224a7d6e09f57dc6278293937d0 Parents: 80fd312 Author: Stefan Seelmann <[email protected]> Authored: Sun Jun 17 09:52:45 2018 +0200 Committer: Bolke de Bruin <[email protected]> Committed: Sat Jun 30 16:19:29 2018 +0200 ---------------------------------------------------------------------- airflow/jobs.py | 1 + .../versions/05f30312d566_merge_heads.py | 43 ++++++++++ .../856955da8476_fix_sqlite_foreign_key.py | 90 ++++++++++++++++++++ ...23433877c24_fix_mysql_not_null_constraint.py | 54 ++++++++++++ airflow/models.py | 22 +++-- airflow/utils/sqlalchemy.py | 6 ++ 6 files changed, 207 insertions(+), 9 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/54da3eaf/airflow/jobs.py ---------------------------------------------------------------------- diff --git a/airflow/jobs.py b/airflow/jobs.py index 5143e8e..70891ab 100644 --- a/airflow/jobs.py +++ b/airflow/jobs.py @@ -96,6 +96,7 @@ class BaseJob(Base, LoggingMixin): __table_args__ = ( Index('job_type_heart', job_type, latest_heartbeat), + Index('idx_job_state_heartbeat', state, latest_heartbeat), ) def __init__( http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/54da3eaf/airflow/migrations/versions/05f30312d566_merge_heads.py ---------------------------------------------------------------------- diff --git a/airflow/migrations/versions/05f30312d566_merge_heads.py b/airflow/migrations/versions/05f30312d566_merge_heads.py new file mode 100644 index 0000000..78d5652 --- /dev/null +++ b/airflow/migrations/versions/05f30312d566_merge_heads.py @@ -0,0 +1,43 @@ +# flake8: noqa +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""merge heads + +Revision ID: 05f30312d566 +Revises: 86770d1215c0, 0e2a74e0fc9f +Create Date: 2018-06-17 10:47:23.339972 + +""" + +# revision identifiers, used by Alembic. +revision = '05f30312d566' +down_revision = ('86770d1215c0', '0e2a74e0fc9f') +branch_labels = None +depends_on = None + +from alembic import op +import sqlalchemy as sa + + +def upgrade(): + pass + + +def downgrade(): + pass http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/54da3eaf/airflow/migrations/versions/856955da8476_fix_sqlite_foreign_key.py ---------------------------------------------------------------------- diff --git a/airflow/migrations/versions/856955da8476_fix_sqlite_foreign_key.py b/airflow/migrations/versions/856955da8476_fix_sqlite_foreign_key.py new file mode 100644 index 0000000..5b11dc7 --- /dev/null +++ b/airflow/migrations/versions/856955da8476_fix_sqlite_foreign_key.py @@ -0,0 +1,90 @@ +# flake8: noqa +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""fix sqlite foreign key + +Revision ID: 856955da8476 +Revises: f23433877c24 +Create Date: 2018-06-17 15:54:53.844230 + +""" + +# revision identifiers, used by Alembic. +revision = '856955da8476' +down_revision = 'f23433877c24' +branch_labels = None +depends_on = None + +from alembic import op +import sqlalchemy as sa + + +def upgrade(): + conn = op.get_bind() + if conn.dialect.name == 'sqlite': + # Fix broken foreign-key constraint for existing SQLite DBs. + # + # Re-define tables and use copy_from to avoid reflection + # which would fail because referenced user table doesn't exist. + # + # Use batch_alter_table to support SQLite workaround. + chart_table = sa.Table('chart', + sa.MetaData(), + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('label', sa.String(length=200), nullable=True), + sa.Column('conn_id', sa.String(length=250), nullable=False), + sa.Column('user_id', sa.Integer(), nullable=True), + sa.Column('chart_type', sa.String(length=100), nullable=True), + sa.Column('sql_layout', sa.String(length=50), nullable=True), + sa.Column('sql', sa.Text(), nullable=True), + sa.Column('y_log_scale', sa.Boolean(), nullable=True), + sa.Column('show_datatable', sa.Boolean(), nullable=True), + sa.Column('show_sql', sa.Boolean(), nullable=True), + sa.Column('height', sa.Integer(), nullable=True), + sa.Column('default_params', sa.String(length=5000), nullable=True), + sa.Column('x_is_date', sa.Boolean(), nullable=True), + sa.Column('iteration_no', sa.Integer(), nullable=True), + sa.Column('last_modified', sa.DateTime(), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + with op.batch_alter_table('chart', copy_from=chart_table) as batch_op: + batch_op.create_foreign_key('chart_user_id_fkey', 'users', + ['user_id'], ['id']) + + known_event_table = sa.Table('known_event', + sa.MetaData(), + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('label', sa.String(length=200), nullable=True), + sa.Column('start_date', sa.DateTime(), nullable=True), + sa.Column('end_date', sa.DateTime(), nullable=True), + sa.Column('user_id', sa.Integer(), nullable=True), + sa.Column('known_event_type_id', sa.Integer(), nullable=True), + sa.Column('description', sa.Text(), nullable=True), + sa.ForeignKeyConstraint(['known_event_type_id'], + ['known_event_type.id'], ), + sa.PrimaryKeyConstraint('id') + ) + with op.batch_alter_table('chart', copy_from=known_event_table) as batch_op: + batch_op.create_foreign_key('known_event_user_id_fkey', 'users', + ['user_id'], ['id']) + + +def downgrade(): + # Downgrade would fail because the broken FK constraint can't be re-created. + pass http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/54da3eaf/airflow/migrations/versions/f23433877c24_fix_mysql_not_null_constraint.py ---------------------------------------------------------------------- diff --git a/airflow/migrations/versions/f23433877c24_fix_mysql_not_null_constraint.py b/airflow/migrations/versions/f23433877c24_fix_mysql_not_null_constraint.py new file mode 100644 index 0000000..44edeef --- /dev/null +++ b/airflow/migrations/versions/f23433877c24_fix_mysql_not_null_constraint.py @@ -0,0 +1,54 @@ +# flake8: noqa +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""fix mysql not null constraint + +Revision ID: f23433877c24 +Revises: 05f30312d566 +Create Date: 2018-06-17 10:16:31.412131 + +""" + +# revision identifiers, used by Alembic. +revision = 'f23433877c24' +down_revision = '05f30312d566' +branch_labels = None +depends_on = None + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql + +def upgrade(): + conn = op.get_bind() + if conn.dialect.name == 'mysql': + conn.execute("SET time_zone = '+00:00'") + op.alter_column('task_fail', 'execution_date', existing_type=mysql.TIMESTAMP(fsp=6), nullable=False) + op.alter_column('xcom', 'execution_date', existing_type=mysql.TIMESTAMP(fsp=6), nullable=False) + op.alter_column('xcom', 'timestamp', existing_type=mysql.TIMESTAMP(fsp=6), nullable=False) + + +def downgrade(): + conn = op.get_bind() + if conn.dialect.name == 'mysql': + conn.execute("SET time_zone = '+00:00'") + op.alter_column('xcom', 'timestamp', existing_type=mysql.TIMESTAMP(fsp=6), nullable=True) + op.alter_column('xcom', 'execution_date', existing_type=mysql.TIMESTAMP(fsp=6), nullable=True) + op.alter_column('task_fail', 'execution_date', existing_type=mysql.TIMESTAMP(fsp=6), nullable=True) + http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/54da3eaf/airflow/models.py ---------------------------------------------------------------------- diff --git a/airflow/models.py b/airflow/models.py index b0841e6..bde6137 100755 --- a/airflow/models.py +++ b/airflow/models.py @@ -56,7 +56,7 @@ from urllib.parse import urlparse, quote, parse_qsl from sqlalchemy import ( Column, Integer, String, DateTime, Text, Boolean, ForeignKey, PickleType, - Index, Float, LargeBinary) + Index, Float, LargeBinary, UniqueConstraint) from sqlalchemy import func, or_, and_, true as sqltrue from sqlalchemy.ext.declarative import declarative_base, declared_attr from sqlalchemy.orm import reconstructor, relationship, synonym @@ -885,7 +885,7 @@ class TaskInstance(Base, LoggingMixin): max_tries = Column(Integer) hostname = Column(String(1000)) unixname = Column(String(1000)) - job_id = Column(Integer, index=True) + job_id = Column(Integer) pool = Column(String(50)) queue = Column(String(50)) priority_weight = Column(Integer) @@ -899,6 +899,7 @@ class TaskInstance(Base, LoggingMixin): Index('ti_state', state), Index('ti_state_lkp', dag_id, task_id, execution_date, state), Index('ti_pool', pool, state, priority_weight), + Index('ti_job_id', job_id), ) def __init__(self, task, execution_date, state=None): @@ -2072,12 +2073,13 @@ class TaskFail(Base): __tablename__ = "task_fail" - task_id = Column(String(ID_LEN), primary_key=True) - dag_id = Column(String(ID_LEN), primary_key=True) - execution_date = Column(UtcDateTime, primary_key=True) + id = Column(Integer, primary_key=True) + task_id = Column(String(ID_LEN), nullable=False) + dag_id = Column(String(ID_LEN), nullable=False) + execution_date = Column(UtcDateTime, nullable=False) start_date = Column(UtcDateTime) end_date = Column(UtcDateTime) - duration = Column(Float) + duration = Column(Integer) __table_args__ = ( Index('idx_task_fail_dag_task_date', dag_id, task_id, execution_date, @@ -4670,8 +4672,8 @@ class DagStat(Base): dag_id = Column(String(ID_LEN), primary_key=True) state = Column(String(50), primary_key=True) - count = Column(Integer, default=0) - dirty = Column(Boolean, default=False) + count = Column(Integer, default=0, nullable=False) + dirty = Column(Boolean, default=False, nullable=False) def __init__(self, dag_id, state, count=0, dirty=False): self.dag_id = dag_id @@ -4804,7 +4806,9 @@ class DagRun(Base, LoggingMixin): dag = None __table_args__ = ( - Index('dr_run_id', dag_id, run_id, unique=True), + Index('dag_id_state', dag_id, _state), + UniqueConstraint('dag_id', 'execution_date'), + UniqueConstraint('dag_id', 'run_id'), ) def __repr__(self): http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/54da3eaf/airflow/utils/sqlalchemy.py ---------------------------------------------------------------------- diff --git a/airflow/utils/sqlalchemy.py b/airflow/utils/sqlalchemy.py index a00fe17..4dab322 100644 --- a/airflow/utils/sqlalchemy.py +++ b/airflow/utils/sqlalchemy.py @@ -101,6 +101,12 @@ def setup_event_handlers( def connect(dbapi_connection, connection_record): connection_record.info['pid'] = os.getpid() + @event.listens_for(engine, "connect") + def set_sqlite_pragma(dbapi_connection, connection_record): + if 'sqlite3.Connection' in str(type(dbapi_connection)): + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() @event.listens_for(engine, "checkout") def checkout(dbapi_connection, connection_record, connection_proxy):
