def upgrade(engine):
session_maker = sessionmaker(bind=engine)
session = session_maker()
r = (session.query(model.ByteHubVersion).order_by(
model.ByteHubVersion.timestamp.desc()).first())
if not r:
current_version = "0"
else:
current_version = r.version
if parse(current_version) >= parse(version.__version__):
# Up-to-date: nothing to do
session.close()
return
if parse(current_version) < parse(version.__version__):
print(f"Upgrading ByteHub database schema...")
with engine.connect() as connection:
context = MigrationContext.configure(connection)
op = Operations(context)
try:
op.add_column("namespace", Column("backend", String(128)))
except OperationalError:
pass
# Add version number
obj = model.ByteHubVersion()
obj.version = version.__version__
session.add(obj)
session.commit()
session.close()
def create_operations(ctx=None, **kwargs):
"""Create an alembic operations object."""
if ctx is None:
ctx = create_migration_ctx(**kwargs)
operations = Operations(ctx)
operations.has_table = has_table
return operations
def upgrade(migrate_engine):
"""
Upgrade operations go here.
Don't create your own engine; bind migrate_engine to your metadata
"""
_reset_base(migrate_engine)
from rhodecode.lib.dbmigrate.schema import db_3_7_0_0
init_model_encryption(db_3_7_0_0)
context = MigrationContext.configure(migrate_engine.connect())
op = Operations(context)
op.create_table(
'external_identities',
sa.Column('provider_name', sa.Unicode(255), primary_key=True),
sa.Column('local_user_id',
sa.Integer(),
sa.ForeignKey('users.user_id'),
primary_key=True),
sa.Column('external_id', sa.Unicode(255), primary_key=True),
sa.Column('external_username', sa.Unicode(1024), default=u''),
sa.Column('access_token', sa.String(1024), default=u''),
sa.Column('alt_token', sa.String(1024), default=u''),
sa.Column('token_secret', sa.String(1024), default=u''))
op.create_index('local_user_id_idx', 'external_identities',
['local_user_id'])
op.create_index('external_id_idx', 'external_identities', ['external_id'])
def create_fk_constraint(self, fk_constraints: list,
const_columns: dict) -> bool:
""" Get list of foreign keys from static list `fk_constraints` and created it """
try:
conn = self.engine.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
for constraint in fk_constraints:
dest_table_name = constraint.pop("table_name")
column_name = constraint.pop("column_name")
source_table = constraint.pop("source_table")
dest_column = constraint.pop("dest_column")
temp = [i for i in const_columns[source_table]]
if not dest_column in temp:
op.create_foreign_key(
None,
source_table,
dest_table_name,
[dest_column],
[column_name],
**constraint,
)
return True
except Exception as err:
logger.error("create_fk_constraint [error] -> %s" % err)
return False
finally:
conn.close()
def add_new_columns(self):
"""
Adds columns present in df but not in the SQL table
for given instance of PandasSpecialEngine.
Notes
-----
Sadly, it seems that we cannot create JSON columns.
"""
# create deepcopies of the column because we are going to unbound
# them for the table model (otherwise alembic would think we add
# a column that already exists in the database)
cols_to_add = [deepcopy(col) for col in self.table.columns
if col.name not in self.get_db_columns_names()]
# check columns are not index levels
if any((c.name in self.df.index.names for c in cols_to_add)):
raise ValueError(('Cannot add any column that is part of the df index!\n'
"You'll have to update your table primary key or change your "
"df index"))
with self.engine.connect() as con:
ctx = MigrationContext.configure(con)
op = Operations(ctx)
for col in cols_to_add:
col.table = None # Important! unbound column from table
op.add_column(self.table.name, col, schema=self.schema)
log(f"Added column {col} (type: {col.type}) in table {self.table.name} "
f'(schema="{self.schema}")')
def setUp(self):
self.conn = config.db.connect()
self.metadata = MetaData()
t1 = Table('foo',
self.metadata,
Column('id', Integer, primary_key=True),
Column('data', String(50)),
Column('x', Integer),
mysql_engine='InnoDB')
t1.create(self.conn)
self.conn.execute(t1.insert(), [{
"id": 1,
"data": "d1",
"x": 5
}, {
"id": 2,
"data": "22",
"x": 6
}, {
"id": 3,
"data": "8.5",
"x": 7
}, {
"id": 4,
"data": "9.46",
"x": 8
}, {
"id": 5,
"data": "d5",
"x": 9
}])
context = MigrationContext.configure(self.conn)
self.op = Operations(context)
def alembic_tests():
print(f'Alembic tests')
conn = engine.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
try:
op.drop_table('waste')
except:
pass
t = op.create_table(
'waste',
Column('bools', sa.Boolean),
Column('ubytes', sa.Tinyint),
Column('shorts', sa.SmallInteger),
Column('ints', sa.Integer),
Column('bigints', sa.BigInteger),
Column('floats', sa.REAL),
Column('doubles', sa.Float),
Column('dates', sa.Date),
Column('datetimes', sa.DateTime),
Column('varchars', sa.String(10)),
Column('nvarchars', sa.UnicodeText),
Column('numerics', sa.Numeric(38, 10)),
)
data = [{
'bools': True,
'ubytes': 5,
'shorts': 55,
'ints': 555,
'bigints': 5555,
'floats': 5.0,
'doubles': 5.5555555,
'dates': date(2012, 11, 23),
'datetimes': datetime(2012, 11, 23, 16, 34, 56),
'varchars': 'bla',
'nvarchars': 'bla2',
'numerics': Decimal("1.1")
}, {
'bools': False,
'ubytes': 6,
'shorts': 66,
'ints': 666,
'bigints': 6666,
'floats': 6.0,
'doubles': 6.6666666,
'dates': date(2012, 11, 24),
'datetimes': datetime(2012, 11, 24, 16, 34, 57),
'varchars': 'bla',
'nvarchars': 'bla2',
'numerics': Decimal("-1.1")
}]
op.bulk_insert(t, data)
res = engine.execute('select * from waste').fetchall()
assert (res == [tuple(dikt.values()) for dikt in data])
def test_alembic_nextbigid_ops(engine):
# Test the migration operations work
with engine.connect() as conn:
context = MigrationContext.configure(conn)
op = Operations(context)
op.create_nextbigid_function()
op.drop_nextbigid_function()
def upgrade(plugin_name=None, revision=None, lower="current"):
database_url = current_app.config.get("SQLALCHEMY_DATABASE_URI")
if database_url.startswith("sqlite"):
current_app.db.create_all()
return
if plugin_name is None:
# Get the directory name of the plugin if unspecified
# Doing it this way doesn't waste the rest of the inspect.stack call
frame = inspect.currentframe()
caller_info = inspect.getframeinfo(frame.f_back)
caller_path = caller_info[0]
plugin_name = os.path.basename(os.path.dirname(caller_path))
# Check if the plugin has migraitons
migrations_path = os.path.join(current_app.plugins_dir, plugin_name,
"migrations")
if os.path.isdir(migrations_path) is False:
return
engine = create_engine(database_url, poolclass=pool.NullPool)
conn = engine.connect()
context = MigrationContext.configure(conn)
op = Operations(context)
# Find the list of migrations to run
config = Config()
config.set_main_option("script_location", migrations_path)
config.set_main_option("version_locations", migrations_path)
script = ScriptDirectory.from_config(config)
# Choose base revision for plugin upgrade
# "current" points to the current plugin version stored in config
# None represents the absolute base layer (e.g. first installation)
if lower == "current":
lower = get_config(plugin_name + "_alembic_version")
# Do we upgrade to head or to a specific revision
if revision is None:
upper = script.get_current_head()
else:
upper = revision
# Apply from lower to upper
revs = list(script.iterate_revisions(lower=lower, upper=upper))
revs.reverse()
try:
for r in revs:
with context.begin_transaction():
r.module.upgrade(op=op)
# Set revision that succeeded so we don't need
# to start from the beginning on failure
set_config(plugin_name + "_alembic_version", r.revision)
finally:
conn.close()
# Set the new latest revision
set_config(plugin_name + "_alembic_version", upper)
def adapt_dtype_of_empty_db_columns(self):
"""
Changes the data types of empty columns in the SQL table defined
in given instance of a PandasSpecialEngine.
This should only happen in case of data type mismatches.
This means with columns for which the sqlalchemy table
model for df and the model for the SQL table have different data types.
"""
empty_db_columns = self.get_empty_columns()
db_table = self.get_db_table_schema()
# if column does not have value in db and there are values
# in the frame then change the column type if needed
for col in empty_db_columns:
# check if the column also exists in df
if col.name not in self.df.columns:
continue
# check same type
orig_type = db_table.columns[col.name].type.compile(self.engine.dialect)
dest_type = self.table.columns[col.name].type.compile(self.engine.dialect)
# remove character count e.g. "VARCHAR(50)" -> "VARCHAR"
orig_type = RE_CHARCOUNT_COL_TYPE.sub('', orig_type)
dest_type = RE_CHARCOUNT_COL_TYPE.sub('', dest_type)
# if same type or we want to insert TEXT instead of JSON continue
# (JSON is not supported on some DBs so it's normal to have TEXT instead)
if ((orig_type == dest_type) or
((orig_type == 'JSON') and (dest_type == 'TEXT'))):
continue
# grab the col/index from the df
# so we can check if there are any values
if col.name in self.df.index.names:
df_col = self.df.index.get_level_values(col.name)
else:
df_col = self.df[col.name]
if df_col.notna().any():
# raise error if we have to modify the dtype but we have a SQlite engine
# (SQLite does not support data type alteration)
if self._db_type == 'sqlite':
raise ValueError('SQlite does not support column data type alteration!')
with self.engine.connect() as con:
ctx = MigrationContext.configure(con)
op = Operations(ctx)
new_col = self.table.columns[col.name]
# check if postgres (in which case we have to use "using" syntax
# to alter columns data types)
if self._db_type == 'postgres':
escaped_col = str(new_col.compile(dialect=self.engine.dialect))
compiled_type = new_col.type.compile(dialect=self.engine.dialect)
alter_kwargs = {'postgresql_using':f'{escaped_col}::{compiled_type}'}
else:
alter_kwargs = {}
op.alter_column(table_name=self.table.name,
column_name=new_col.name,
type_=new_col.type,
schema=self.schema,
**alter_kwargs)
log(f"Changed type of column {new_col.name} "
f"from {col.type} to {new_col.type} "
f'in table {self.table.name} (schema="{self.schema}")')
def _make_alembic_op(self, conn):
'''
Return an instance Operations which is used to
generate DDL statement
'''
migration_ctx = MigrationContext.configure(conn)
op = Operations(migration_ctx)
return op
def get_upgrade_op(session):
"""
Create a migration context and an operations object for performing upgrades.
:param session: The SQLAlchemy session object.
"""
context = MigrationContext.configure(session.bind.connect())
return Operations(context)
def test_standalone_op():
eng, buf = capture_db()
env = MigrationContext.configure(eng)
op = Operations(env)
op.alter_column("t", "c", nullable=True)
eq_(buf, ['ALTER TABLE t ALTER COLUMN c DROP NOT NULL'])
def alter_add_foreign_key_to_backtest_summary(self):
ctx = MigrationContext.configure(self.db_client.connector)
op = Operations(ctx)
with op.batch_alter_table(
self.backtest_management_table_name) as batch_op:
batch_op.create_foreign_key("fk_management_summary",
"backtest_summary",
["backtest_summary_id"], ["id"])
def upgrade():
dburl = "mysql+mysqlconnector://user:password@localhost/dbname?charset=utf8"
syncinst = SyncDB(Base, dburl)
engine = syncinst.getengine()
conn = engine.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
op.add_column('nt_plfbgcsjb',
Column('FILECOUNT', TINYINT, nullable=True, default=0))
def _fixture(self):
migration_context = mock.Mock(opts={})
op = Operations(migration_context)
batch = op.batch_alter_table('tname', recreate='never').__enter__()
with mock.patch("alembic.operations.sa_schema") as mock_schema:
yield batch
batch.impl.flush()
self.mock_schema = mock_schema
def init_online(app: Flask):
db: SQLAlchemy = app.extensions['sqlalchemy'].db
conn = db.engine.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
try:
op.drop_table('seeds')
except:
pass
op.create_table('seeds', sa.Column('file', sa.String(255), primary_key=True))
def fix(self):
context = MigrationContext.configure(self.parent.engine.connect())
op = Operations(context)
for table in self.parent.base.metadata.sorted_tables:
if table.name == self.table:
for column in table.columns:
if column.name == self.name:
with op.batch_alter_table(table.name) as batch_op:
batch_op.add_column(column.copy())
return
def update_ego_grid_ehv_substation(**kwargs):
postgres_hook = PostgresHook(postgres_conn_id='postgres_oedb')
engine = postgres_hook.get_sqlalchemy_engine()
conn = engine.connect()
#alembic
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
#add column
op.add_column('ego_grid_ehv_substation',
Column('otg_id', BigInteger),
schema='model_draft')
# load the tables
meta = MetaData()
conn.execute("SET search_path TO model_draft, grid, public")
ego_grid_hvmv_substation = Table('ego_grid_hvmv_substation',
meta,
autoload=True,
autoload_with=conn,
postgresql_ignore_search_path=True)
ego_grid_ehv_substation = Table('ego_grid_ehv_substation',
meta,
autoload=True,
autoload_with=conn,
postgresql_ignore_search_path=True)
otg_ehvhv_bus_data = Table('otg_ehvhv_bus_data',
meta,
autoload=True,
autoload_with=conn,
postgresql_ignore_search_path=True)
#operations with "ego_grid_ehv_substation" table
#update
ehv_substation_update = ego_grid_ehv_substation.update().values(
otg_id=otg_ehvhv_bus_data.c.bus_i).where(
and_(
otg_ehvhv_bus_data.c.base_kv > 110,
otg_ehvhv_bus_data.c.osm_substation_id == cast(
func.trim(ego_grid_ehv_substation.c.osm_id, 'nwr'),
BigInteger)))
#delete
ehv_substation_delete = ego_grid_ehv_substation.delete().where(
ego_grid_ehv_substation.c.otg_id.is_(None))
#execution
conn.execute(ehv_substation_update)
conn.execute(ehv_substation_delete)
def change_column_nullable(table_oid, column_index, nullable, engine):
table = tables.reflect_table_from_oid(table_oid, engine)
column = table.columns[column_index]
with engine.begin() as conn:
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
op.alter_column(table.name,
column.name,
nullable=nullable,
schema=table.schema)
return tables.reflect_table_from_oid(table_oid,
engine).columns[column_index]
def drop_fk(self, fk_constraints: str):
try:
conn = self.engine.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
# [('todo', 'todo_ibfk_1')]
for fk in fk_constraints:
op.drop_constraint(fk[1], fk[0], type_="foreignkey")
except Exception as err:
logger.error("fk_drop [error] -> %s" % err)
finally:
conn.close()
def update(self):
code = self.migration_code
migrate = _get_migrate_function(code)
schema_name = self.context.content_metadata().schema
with self.engine.connect() as conn:
migration_context = MigrationContext.configure(
conn, opts={"include_schemas": False})
op = Operations(migration_context)
if not self.engine.dialect.has_schema(self.engine, schema_name):
conn.execute(CreateSchema(schema_name))
migrate(op)
def drop_column(
engine,
table_oid,
column_index,
):
column_index = int(column_index)
table = tables.reflect_table_from_oid(table_oid, engine)
column = table.columns[column_index]
with engine.begin() as conn:
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
op.drop_column(table.name, column.name, schema=table.schema)
def add_columns(engine, raw_diffs, table_names):
with engine.begin() as conn:
ctx = get_migration_context(conn, table_names)
op = Operations(ctx)
columns = _get_columns_to_add(raw_diffs, table_names)
for col in columns:
table_name = col.table.name
# the column has a reference to a table definition that already
# has the column defined, so remove that and add the column
col.table = None
op.add_column(table_name, col)
def _fixture(self):
self.metadata = MetaData()
self.table = Table(
'foo', self.metadata,
Column('id', Integer, primary_key=True),
Column('data', String(50)),
Column('x', Integer),
)
context = op_fixture(dialect="sqlite", as_sql=True)
self.op = Operations(context)
return context
def _fixture(self, schema=None):
migration_context = mock.Mock(
opts={}, impl=mock.MagicMock(__dialect__='sqlite'))
op = Operations(migration_context)
batch = op.batch_alter_table('tname', recreate='never',
schema=schema).__enter__()
mock_schema = mock.MagicMock()
with mock.patch("alembic.operations.schemaobj.sa_schema", mock_schema):
yield batch
batch.impl.flush()
self.mock_schema = mock_schema
def update_column(self, table_name, column_name, col_type, **options):
""" Updated existing table column with new column """
try:
conn = self.engine.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
# op.alter_column(table_name, column_name,type_=col_type,postgresql_using=f"{column_name}::{col_type}") #FIXME not working
except Exception as err:
logger.error("update_column [error] -> %s" % err)
finally:
conn.close()
def _alter_column_name(self, base, old_name, new_name):
# NOTE: Create alembic connection and operation object! By John Doe
db_conn = config.ENGINE.connect()
alembic_ctx = MigrationContext.configure(db_conn)
alembic_op = Operations(alembic_ctx)
doc_table = get_doc_table(base.metadata.name, config.METADATA,
**base.relational_fields)
alembic_op.alter_column(doc_table.name,
old_name,
new_column_name=new_name)
db_conn.close()
def get_db_tables(conn):
"""Get current and default table values from the db.
:param engine: Initialized alembic engine object.
:type engine: object
:returns: tuple
"""
query = text("SELECT TABLE_NAME from information_schema.tables\
WHERE TABLE_NAME\
LIKE '%alembic_version%'\
AND table_schema = 'refstack'")
context = alembic_migration.MigrationContext.configure(conn)
op = Operations(context)
connection = op.get_bind()
search = connection.execute(query)
result = search.fetchall()
if isinstance(result, Iterable):
result = [table[0] for table in result]
else:
result = None
# if there is more than one version table, modify the
# one that does not have the default name, because subunit2sql uses the
# default name.
if result:
current_name =\
next((table for table in result if table != "alembic_version"),
result[0])
current_name = current_name.decode('utf-8')
current_version = get_table_version(conn, current_name)
default_name =\
next((table for table in result
if table == "alembic_version"), None)
default_version = get_table_version(conn, default_name)
if len(result) > 1 and not current_version:
if not default_name:
# this is the case where there is more than one
# nonstandard-named alembic table, and no default
current_name = next(
(table for table in result if table != current_name),
result[0])
current_name = current_name.decode('utf-8')
elif current_name:
# this is the case where the current-named table
# exists, but is empty
current_name = default_name
current_version = default_version
current_table = (current_name, current_version)
default_table = (default_name, default_version)
else:
default_table = (None, None)
current_table = default_table
return current_table, default_table
def fix(self):
"""Uses Alembic batch operations to alter the column datatype in the table.
"""
context = MigrationContext.configure(self.parent.engine.connect())
op = Operations(context)
for table in self.parent.base.metadata.sorted_tables:
if table.name == self.table:
for column in table.columns:
if column.name == self.column['name']:
with op.batch_alter_table(table.name) as batch_op:
batch_op.alter_column(column.name,
type_=column.type)
return
