def upgrade(): bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) with op.batch_alter_table("tab_state") as batch_op: table_schema_id_constraint = generic_find_fk_constraint_name( "tab_state", {"id"}, "dbs", insp) if table_schema_id_constraint: batch_op.drop_constraint( table_schema_id_constraint, type_="foreignkey", ) table_schema_id_constraint = generic_find_fk_constraint_name( "tab_state", {"client_id"}, "query", insp) if table_schema_id_constraint: batch_op.drop_constraint( table_schema_id_constraint, type_="foreignkey", ) batch_op.create_foreign_key( "tab_state_database_id_fkey", "dbs", ["database_id"], ["id"], ondelete="CASCADE", ) batch_op.create_foreign_key( "tab_state_latest_query_id_fkey", "query", ["latest_query_id"], ["client_id"], ondelete="SET NULL", ) with op.batch_alter_table("table_schema") as batch_op: table_schema_id_constraint = generic_find_fk_constraint_name( "table_schema", {"id"}, "dbs", insp) if table_schema_id_constraint: batch_op.drop_constraint( table_schema_id_constraint, type_="foreignkey", ) batch_op.create_foreign_key( "table_schema_database_id_fkey", "dbs", ["database_id"], ["id"], ondelete="CASCADE", )
def upgrade(): bind = op.get_bind() metadata = sa.MetaData(bind=bind) insp = sa.engine.reflection.Inspector.from_engine(bind) rls_filter_tables = op.create_table( "rls_filter_tables", sa.Column("id", sa.Integer(), nullable=False), sa.Column("table_id", sa.Integer(), nullable=True), sa.Column("rls_filter_id", sa.Integer(), nullable=True), sa.ForeignKeyConstraint(["rls_filter_id"], ["row_level_security_filters.id"]), sa.ForeignKeyConstraint(["table_id"], ["tables.id"]), sa.PrimaryKeyConstraint("id"), ) rlsf = sa.Table("row_level_security_filters", metadata, autoload=True) filter_ids = sa.select([rlsf.c.id, rlsf.c.table_id]) for row in bind.execute(filter_ids): move_table_id = rls_filter_tables.insert().values( rls_filter_id=row["id"], table_id=row["table_id"]) bind.execute(move_table_id) with op.batch_alter_table("row_level_security_filters") as batch_op: fk_constraint_name = generic_find_fk_constraint_name( "row_level_security_filters", {"id"}, "tables", insp) if fk_constraint_name: batch_op.drop_constraint(fk_constraint_name, type_="foreignkey") batch_op.drop_column("table_id")
def upgrade(): op.create_table('sqlatable_user', sa.Column('id', sa.Integer(), nullable=False), sa.Column('user_id', sa.Integer(), nullable=True), sa.Column('table_id', sa.Integer(), nullable=True), sa.ForeignKeyConstraint( ['table_id'], ['tables.id'], ), sa.ForeignKeyConstraint( ['user_id'], ['ab_user.id'], ), sa.PrimaryKeyConstraint('id')) op.create_table( 'druiddatasource_user', sa.Column('id', sa.Integer(), nullable=False), sa.Column('user_id', sa.Integer(), nullable=True), sa.Column('datasource_id', sa.Integer(), nullable=True), sa.ForeignKeyConstraint( ['datasource_id'], ['datasources.id'], ), sa.ForeignKeyConstraint( ['user_id'], ['ab_user.id'], ), sa.PrimaryKeyConstraint('id')) bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) session = db.Session(bind=bind) tables = session.query(SqlaTable).all() for table in tables: if table.user_id is not None: session.execute(sqlatable_user.insert().values( user_id=table.user_id, table_id=table.id)) druiddatasources = session.query(DruidDatasource).all() for druiddatasource in druiddatasources: if druiddatasource.user_id is not None: session.execute(druiddatasource_user.insert().values( user_id=druiddatasource.user_id, datasource_id=druiddatasource.id)) session.close() with op.batch_alter_table('tables') as batch_op: batch_op.drop_constraint('user_id', type_='foreignkey') batch_op.drop_column('user_id') with op.batch_alter_table('datasources') as batch_op: batch_op.drop_constraint(generic_find_fk_constraint_name( 'datasources', {'id'}, 'ab_user', insp, ), type_='foreignkey') batch_op.drop_column('user_id')
def upgrade(): op.create_table( "sqlatable_user", sa.Column("id", sa.Integer(), nullable=False), sa.Column("user_id", sa.Integer(), nullable=True), sa.Column("table_id", sa.Integer(), nullable=True), sa.ForeignKeyConstraint(["table_id"], ["tables.id"]), sa.ForeignKeyConstraint(["user_id"], ["ab_user.id"]), sa.PrimaryKeyConstraint("id"), ) op.create_table( "druiddatasource_user", sa.Column("id", sa.Integer(), nullable=False), sa.Column("user_id", sa.Integer(), nullable=True), sa.Column("datasource_id", sa.Integer(), nullable=True), sa.ForeignKeyConstraint(["datasource_id"], ["datasources.id"]), sa.ForeignKeyConstraint(["user_id"], ["ab_user.id"]), sa.PrimaryKeyConstraint("id"), ) bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) session = db.Session(bind=bind) tables = session.query(SqlaTable).all() for table in tables: if table.user_id is not None: session.execute( sqlatable_user.insert().values(user_id=table.user_id, table_id=table.id) ) druiddatasources = session.query(DruidDatasource).all() for druiddatasource in druiddatasources: if druiddatasource.user_id is not None: session.execute( druiddatasource_user.insert().values( user_id=druiddatasource.user_id, datasource_id=druiddatasource.id ) ) session.close() with op.batch_alter_table("tables") as batch_op: batch_op.drop_constraint("user_id", type_="foreignkey") batch_op.drop_column("user_id") with op.batch_alter_table("datasources") as batch_op: batch_op.drop_constraint( generic_find_fk_constraint_name("datasources", {"id"}, "ab_user", insp), type_="foreignkey", ) batch_op.drop_column("user_id")
def upgrade(): op.create_table('sqlatable_user', sa.Column('id', sa.Integer(), nullable=False), sa.Column('user_id', sa.Integer(), nullable=True), sa.Column('table_id', sa.Integer(), nullable=True), sa.ForeignKeyConstraint(['table_id'], ['tables.id'], ), sa.ForeignKeyConstraint(['user_id'], ['ab_user.id'], ), sa.PrimaryKeyConstraint('id') ) op.create_table('druiddatasource_user', sa.Column('id', sa.Integer(), nullable=False), sa.Column('user_id', sa.Integer(), nullable=True), sa.Column('datasource_id', sa.Integer(), nullable=True), sa.ForeignKeyConstraint(['datasource_id'], ['datasources.id'], ), sa.ForeignKeyConstraint(['user_id'], ['ab_user.id'], ), sa.PrimaryKeyConstraint('id') ) bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) session = db.Session(bind=bind) tables = session.query(SqlaTable).all() for table in tables: if table.user_id is not None: session.execute( sqlatable_user.insert().values(user_id=table.user_id, table_id=table.id) ) druiddatasources = session.query(DruidDatasource).all() for druiddatasource in druiddatasources: if druiddatasource.user_id is not None: session.execute( druiddatasource_user.insert().values(user_id=druiddatasource.user_id, datasource_id=druiddatasource.id) ) session.close() with op.batch_alter_table('tables') as batch_op: batch_op.drop_constraint('user_id', type_='foreignkey') batch_op.drop_column('user_id') with op.batch_alter_table('datasources') as batch_op: batch_op.drop_constraint(generic_find_fk_constraint_name( 'datasources', {'id'}, 'ab_user', insp, ), type_='foreignkey') batch_op.drop_column('user_id')
def downgrade(): bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) with op.batch_alter_table("tab_state") as batch_op: batch_op.drop_constraint( generic_find_fk_constraint_name("tab_state", {"id"}, "saved_query", insp), type_="foreignkey", ) batch_op.create_foreign_key( "saved_query_id", "saved_query", ["saved_query_id"], ["id"], )
def downgrade(): bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) # First, drop the foreign key constraint prior to altering columns. fk_datasources_cluster_name_clusters = (generic_find_fk_constraint_name( "datasources", {"cluster_name"}, "clusters", insp) or "fk_datasources_cluster_name_clusters") with op.batch_alter_table("datasources", naming_convention=conv) as batch_op: batch_op.drop_constraint(fk_datasources_cluster_name_clusters, type_="foreignkey") # Second, make the columns nullable. with op.batch_alter_table("datasources") as batch_op: batch_op.alter_column("cluster_name", existing_type=sa.String(250), nullable=True) with op.batch_alter_table("clusters") as batch_op: batch_op.alter_column("cluster_name", existing_type=sa.String(250), nullable=True) with op.batch_alter_table("dbs") as batch_op: batch_op.alter_column("database_name", existing_type=sa.String(250), nullable=True) with op.batch_alter_table("tables") as batch_op: batch_op.alter_column("table_name", existing_type=sa.String(250), nullable=True) # Finally, re-add the foreign key constraint. with op.batch_alter_table("datasources") as batch_op: batch_op.create_foreign_key( fk_datasources_cluster_name_clusters, "clusters", ["cluster_name"], ["cluster_name"], )
def downgrade(): bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) # Add cluster_name column with op.batch_alter_table("datasources") as batch_op: batch_op.add_column(sa.Column("cluster_name", sa.String(250))) # Update cluster_name values metadata = sa.MetaData(bind=bind) datasources = sa.Table("datasources", metadata, autoload=True) clusters = sa.Table("clusters", metadata, autoload=True) statement = datasources.update().values( cluster_name=sa.select([clusters.c.cluster_name]).where( datasources.c.cluster_id == clusters.c.id).as_scalar()) bind.execute(statement) with op.batch_alter_table("datasources") as batch_op: # Drop cluster_id column fk_constraint_name = generic_find_fk_constraint_name( "datasources", {"id"}, "clusters", insp) uq_constraint_name = generic_find_uq_constraint_name( "datasources", {"cluster_id", "datasource_name"}, insp) batch_op.drop_constraint(fk_constraint_name, type_="foreignkey") batch_op.drop_constraint(uq_constraint_name, type_="unique") batch_op.drop_column("cluster_id") # Add constraints to cluster_name column batch_op.alter_column("cluster_name", existing_type=sa.String(250), nullable=False) batch_op.create_unique_constraint("uq_datasources_cluster_name", ["cluster_name", "datasource_name"]) batch_op.create_foreign_key( "fk_datasources_cluster_name_clusters", "clusters", ["cluster_name"], ["cluster_name"], )
def downgrade(): bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) # Add the new more restrictive uniqueness constraint which is required by # the foreign key constraints. Note this operation will fail if the # datasources.datasource_name column is no longer unique. with op.batch_alter_table("datasources", naming_convention=conv) as batch_op: batch_op.create_unique_constraint("uq_datasources_datasource_name", ["datasource_name"]) # Augment the tables which have a foreign key constraint related to the # datasources.datasource_id column. for foreign in ["columns", "metrics"]: with op.batch_alter_table(foreign, naming_convention=conv) as batch_op: # Add the datasource_name column with the relevant constraints. batch_op.add_column(sa.Column("datasource_name", sa.String(255))) batch_op.create_foreign_key( "fk_{}_datasource_name_datasources".format(foreign), "datasources", ["datasource_name"], ["datasource_name"], ) # Helper table for database migration using minimal schema. table = sa.Table( foreign, sa.MetaData(), sa.Column("id", sa.Integer, primary_key=True), sa.Column("datasource_name", sa.String(255)), sa.Column("datasource_id", sa.Integer), ) # Migrate the existing data. for datasource in bind.execute(datasources.select()): bind.execute(table.update().where( table.c.datasource_id == datasource.id).values( datasource_name=datasource.datasource_name)) with op.batch_alter_table(foreign, naming_convention=conv) as batch_op: # Drop the datasource_id column and associated constraint. batch_op.drop_constraint( "fk_{}_datasource_id_datasources".format(foreign), type_="foreignkey") batch_op.drop_column("datasource_id") with op.batch_alter_table("datasources", naming_convention=conv) as batch_op: # Prior to dropping the uniqueness constraint, the foreign key # associated with the cluster_name column needs to be dropped. batch_op.drop_constraint( generic_find_fk_constraint_name("datasources", {"cluster_name"}, "clusters", insp) or "fk_datasources_cluster_name_clusters", type_="foreignkey", ) # Drop the old less restrictive uniqueness constraint. batch_op.drop_constraint( generic_find_uq_constraint_name( "datasources", {"cluster_name", "datasource_name"}, insp) or "uq_datasources_cluster_name", type_="unique", ) # Re-create the foreign key associated with the cluster_name column. batch_op.create_foreign_key( "fk_{}_datasource_id_datasources".format(foreign), "clusters", ["cluster_name"], ["cluster_name"], )
def upgrade(): bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) if isinstance(bind.dialect, SQLiteDialect): op.add_column( "alerts", sa.Column("validator_config", sa.Text(), server_default="", nullable=True), ) op.add_column( "alerts", sa.Column("database_id", sa.Integer(), server_default="0", nullable=False), ) op.add_column( "alerts", sa.Column("sql", sa.Text(), server_default="", nullable=False) ) op.add_column( "alerts", sa.Column( "validator_type", sa.String(length=100), server_default="", nullable=False, ), ) else: # mysql does not support server_default for text fields op.add_column( "alerts", sa.Column("validator_config", sa.Text(), default="", nullable=True), ) op.add_column( "alerts", sa.Column("database_id", sa.Integer(), default=0, nullable=False), ) op.add_column("alerts", sa.Column("sql", sa.Text(), default="", nullable=False)) op.add_column( "alerts", sa.Column( "validator_type", sa.String(length=100), default="", nullable=False ), ) # Migrate data session = db.Session(bind=bind) alerts = session.query(Alert).all() for a in alerts: if a.sql_observer: a.sql = a.sql_observer[0].sql a.database_id = a.sql_observer[0].database_id if a.validators: a.validator_type = a.validators[0].validator_type a.validator_config = a.validators[0].config session.commit() if not isinstance(bind.dialect, SQLiteDialect): constraint = generic_find_fk_constraint_name( "sql_observations", {"id"}, "sql_observers", insp ) op.drop_constraint(constraint, "sql_observations", type_="foreignkey") op.drop_column("sql_observations", "observer_id") op.drop_table("alert_validators") op.drop_table("sql_observers") # sqlite does not support column and fk deletion if isinstance(bind.dialect, SQLiteDialect): op.drop_table("sql_observations") op.create_table( "sql_observations", sa.Column("id", sa.Integer(), nullable=False), sa.Column("dttm", sa.DateTime(), nullable=True), sa.Column("alert_id", sa.Integer(), nullable=True), sa.Column("value", sa.Float(), nullable=True), sa.Column("error_msg", sa.String(length=500), nullable=True), sa.ForeignKeyConstraint(["alert_id"], ["alerts.id"],), sa.PrimaryKeyConstraint("id"), ) else: op.create_foreign_key(None, "alerts", "dbs", ["database_id"], ["id"])
def downgrade(): bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) op.create_table( "sql_observers", sa.Column("created_on", sa.DateTime(), nullable=True), sa.Column("changed_on", sa.DateTime(), nullable=True), sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), sa.Column("sql", sa.Text(), nullable=False), sa.Column("created_by_fk", sa.Integer(), autoincrement=False, nullable=True), sa.Column("changed_by_fk", sa.Integer(), autoincrement=False, nullable=True), sa.Column("alert_id", sa.Integer(), autoincrement=False, nullable=False), sa.Column("database_id", sa.Integer(), autoincrement=False, nullable=False), sa.ForeignKeyConstraint(["alert_id"], ["alerts.id"]), sa.ForeignKeyConstraint(["changed_by_fk"], ["ab_user.id"]), sa.ForeignKeyConstraint(["created_by_fk"], ["ab_user.id"]), sa.ForeignKeyConstraint(["database_id"], ["dbs.id"]), sa.PrimaryKeyConstraint("id"), ) op.create_table( "alert_validators", sa.Column("created_on", sa.DateTime(), nullable=True), sa.Column("changed_on", sa.DateTime(), nullable=True), sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), sa.Column("validator_type", sa.String(length=100), nullable=False,), sa.Column("config", sa.Text(), nullable=True), sa.Column("created_by_fk", sa.Integer(), autoincrement=False, nullable=True), sa.Column("changed_by_fk", sa.Integer(), autoincrement=False, nullable=True), sa.Column("alert_id", sa.Integer(), autoincrement=False, nullable=False), sa.ForeignKeyConstraint( ["alert_id"], ["alerts.id"], name="alert_validators_ibfk_1" ), sa.ForeignKeyConstraint( ["changed_by_fk"], ["ab_user.id"], name="alert_validators_ibfk_2" ), sa.ForeignKeyConstraint( ["created_by_fk"], ["ab_user.id"], name="alert_validators_ibfk_3" ), sa.PrimaryKeyConstraint("id"), ) # Migrate data session = db.Session(bind=bind) alerts = session.query(Alert).all() for a in alerts: if a.sql: ob = SQLObserver(sql=a.sql, database_id=a.database_id) a.sql_observer.append(ob) session.add(ob) if a.validator_type: val = Validator( validator_type=a.validator_type, config=a.validator_config, alert_id=a.id, ) a.validators.append(val) session.add(val) session.commit() # sqlite does not support dropping columns if isinstance(bind.dialect, SQLiteDialect): op.add_column( "sql_observations", sa.Column( "observer_id", sa.Integer(), autoincrement=False, nullable=False, server_default="0", ), ) op.drop_table("alerts") op.create_table( "alerts", sa.Column("id", sa.Integer(), nullable=False), sa.Column("label", sa.String(length=150), nullable=False), sa.Column("active", sa.Boolean(), nullable=True), sa.Column("crontab", sa.String(length=50), nullable=False), sa.Column("alert_type", sa.String(length=50), nullable=True), sa.Column("log_retention", sa.Integer(), nullable=False, default=90), sa.Column( "grace_period", sa.Integer(), nullable=False, default=60 * 60 * 24 ), sa.Column("recipients", sa.Text(), nullable=True), sa.Column("slice_id", sa.Integer(), nullable=True), sa.Column("dashboard_id", sa.Integer(), nullable=True), sa.Column("last_eval_dttm", sa.DateTime(), nullable=True), sa.Column("last_state", sa.String(length=10), nullable=True), sa.Column("changed_by_fk", sa.Integer(), nullable=True), sa.Column("changed_on", sa.DateTime(), nullable=True), sa.Column("created_by_fk", sa.Integer(), nullable=True), sa.Column("created_on", sa.DateTime(), nullable=True), sa.Column("slack_channel", sa.Text(), nullable=True), sa.ForeignKeyConstraint(["dashboard_id"], ["dashboards.id"],), sa.ForeignKeyConstraint(["slice_id"], ["slices.id"],), sa.ForeignKeyConstraint(["created_by_fk"], ["ab_user.id"],), sa.ForeignKeyConstraint(["changed_by_fk"], ["ab_user.id"],), sa.PrimaryKeyConstraint("id"), ) else: op.add_column( "sql_observations", sa.Column( "observer_id", sa.Integer(), autoincrement=False, nullable=False, default=0, ), ) constraint = generic_find_fk_constraint_name("alerts", {"id"}, "dbs", insp) op.drop_constraint(constraint, "alerts", type_="foreignkey") op.drop_column("alerts", "validator_type") op.drop_column("alerts", "sql") op.drop_column("alerts", "database_id") op.drop_column("alerts", "validator_config") op.create_foreign_key( "sql_observations_ibfk_2", "sql_observations", "sql_observers", ["observer_id"], ["id"], )
def downgrade(): bind = op.get_bind() insp = sa.engine.reflection.Inspector.from_engine(bind) # Add the new more restrictive uniqueness constraint which is required by # the foreign key constraints. Note this operation will fail if the # datasources.datasource_name column is no longer unique. with op.batch_alter_table('datasources', naming_convention=conv) as batch_op: batch_op.create_unique_constraint( 'uq_datasources_datasource_name', ['datasource_name'], ) # Augment the tables which have a foreign key constraint related to the # datasources.datasource_id column. for foreign in ['columns', 'metrics']: with op.batch_alter_table(foreign, naming_convention=conv) as batch_op: # Add the datasource_name column with the relevant constraints. batch_op.add_column(sa.Column('datasource_name', sa.String(255))) batch_op.create_foreign_key( 'fk_{}_datasource_name_datasources'.format(foreign), 'datasources', ['datasource_name'], ['datasource_name'], ) # Helper table for database migration using minimal schema. table = sa.Table( foreign, sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True), sa.Column('datasource_name', sa.String(255)), sa.Column('datasource_id', sa.Integer), ) # Migrate the existing data. for datasource in bind.execute(datasources.select()): bind.execute( table.update().where( table.c.datasource_id == datasource.id, ).values( datasource_name=datasource.datasource_name, ), ) with op.batch_alter_table(foreign, naming_convention=conv) as batch_op: # Drop the datasource_id column and associated constraint. batch_op.drop_constraint( 'fk_{}_datasource_id_datasources'.format(foreign), type_='foreignkey', ) batch_op.drop_column('datasource_id') with op.batch_alter_table('datasources', naming_convention=conv) as batch_op: # Prior to dropping the uniqueness constraint, the foreign key # associated with the cluster_name column needs to be dropped. batch_op.drop_constraint( generic_find_fk_constraint_name( 'datasources', {'cluster_name'}, 'clusters', insp, ) or 'fk_datasources_cluster_name_clusters', type_='foreignkey', ) # Drop the old less restrictive uniqueness constraint. batch_op.drop_constraint( generic_find_uq_constraint_name( 'datasources', {'cluster_name', 'datasource_name'}, insp, ) or 'uq_datasources_cluster_name', type_='unique', ) # Re-create the foreign key associated with the cluster_name column. batch_op.create_foreign_key( 'fk_{}_datasource_id_datasources'.format(foreign), 'clusters', ['cluster_name'], ['cluster_name'], )