173 lines
6.0 KiB
Python
173 lines
6.0 KiB
Python
# Copyright 2014 Mirantis Inc.
|
|
# Copyright 2014 OpenStack Foundation
|
|
#
|
|
# Licensed 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.
|
|
|
|
import six
|
|
import uuid
|
|
|
|
from alembic import op
|
|
import sqlalchemy
|
|
|
|
|
|
def sqlite_alter_columns(table_name, column_defs):
|
|
"""Implement alter columns for SQLite.
|
|
|
|
The ALTER COLUMN command isn't supported by SQLite specification.
|
|
Instead of calling ALTER COLUMN it uses the following workaround:
|
|
|
|
* create temp table '{table_name}_{rand_uuid}', with some column
|
|
defs replaced;
|
|
* copy all data to the temp table;
|
|
* drop old table;
|
|
* rename temp table to the old table name.
|
|
"""
|
|
connection = op.get_bind()
|
|
meta = sqlalchemy.MetaData(bind=connection)
|
|
meta.reflect()
|
|
|
|
changed_columns = {}
|
|
indexes = []
|
|
for col in column_defs:
|
|
# If we are to have an index on the column, don't create it
|
|
# immediately, instead, add it to a list of indexes to create
|
|
# after the table rename.
|
|
if col.index:
|
|
indexes.append(('ix_%s_%s' % (table_name, col.name),
|
|
table_name,
|
|
[col.name],
|
|
col.unique))
|
|
col.unique = False
|
|
col.index = False
|
|
changed_columns[col.name] = col
|
|
|
|
# construct lists of all columns and their names
|
|
old_columns = []
|
|
new_columns = []
|
|
column_names = []
|
|
for column in meta.tables[table_name].columns:
|
|
column_names.append(column.name)
|
|
old_columns.append(column)
|
|
if column.name in changed_columns.keys():
|
|
new_columns.append(changed_columns[column.name])
|
|
else:
|
|
col_copy = column.copy()
|
|
new_columns.append(col_copy)
|
|
|
|
for key in meta.tables[table_name].foreign_keys:
|
|
constraint = key.constraint
|
|
con_copy = constraint.copy()
|
|
new_columns.append(con_copy)
|
|
|
|
for index in meta.tables[table_name].indexes:
|
|
# If this is a single column index for a changed column, don't
|
|
# copy it because we may already be creating a new version of
|
|
# it (or removing it).
|
|
idx_columns = [col.name for col in index.columns]
|
|
if len(idx_columns)==1 and idx_columns[0] in changed_columns.keys():
|
|
continue
|
|
# Otherwise, recreate the index.
|
|
indexes.append((index.name,
|
|
table_name,
|
|
[col.name for col in index.columns],
|
|
index.unique))
|
|
|
|
# create temp table
|
|
tmp_table_name = "%s_%s" % (table_name, six.text_type(uuid.uuid4()))
|
|
op.create_table(tmp_table_name, *new_columns)
|
|
meta.reflect()
|
|
|
|
try:
|
|
# copy data from the old table to the temp one
|
|
sql_select = sqlalchemy.sql.select(old_columns)
|
|
connection.execute(sqlalchemy.sql.insert(meta.tables[tmp_table_name])
|
|
.from_select(column_names, sql_select))
|
|
except Exception:
|
|
op.drop_table(tmp_table_name)
|
|
raise
|
|
|
|
# drop the old table and rename temp table to the old table name
|
|
op.drop_table(table_name)
|
|
op.rename_table(tmp_table_name, table_name)
|
|
|
|
# (re-)create indexes
|
|
for index in indexes:
|
|
op.create_index(op.f(index[0]), index[1], index[2], unique=index[3])
|
|
|
|
def sqlite_drop_columns(table_name, drop_columns):
|
|
"""Implement drop columns for SQLite.
|
|
|
|
The DROP COLUMN command isn't supported by SQLite specification.
|
|
Instead of calling DROP COLUMN it uses the following workaround:
|
|
|
|
* create temp table '{table_name}_{rand_uuid}', without
|
|
dropped columns;
|
|
* copy all data to the temp table;
|
|
* drop old table;
|
|
* rename temp table to the old table name.
|
|
"""
|
|
connection = op.get_bind()
|
|
meta = sqlalchemy.MetaData(bind=connection)
|
|
meta.reflect()
|
|
|
|
# construct lists of all columns and their names
|
|
old_columns = []
|
|
new_columns = []
|
|
column_names = []
|
|
indexes = []
|
|
for column in meta.tables[table_name].columns:
|
|
if column.name not in drop_columns:
|
|
old_columns.append(column)
|
|
column_names.append(column.name)
|
|
col_copy = column.copy()
|
|
new_columns.append(col_copy)
|
|
|
|
for key in meta.tables[table_name].foreign_keys:
|
|
constraint = key.constraint
|
|
con_copy = constraint.copy()
|
|
new_columns.append(con_copy)
|
|
|
|
for index in meta.tables[table_name].indexes:
|
|
# If this is a single column index for a dropped column, don't
|
|
# copy it.
|
|
idx_columns = [col.name for col in index.columns]
|
|
if len(idx_columns)==1 and idx_columns[0] in drop_columns:
|
|
continue
|
|
# Otherwise, recreate the index.
|
|
indexes.append((index.name,
|
|
table_name,
|
|
[col.name for col in index.columns],
|
|
index.unique))
|
|
|
|
# create temp table
|
|
tmp_table_name = "%s_%s" % (table_name, six.text_type(uuid.uuid4()))
|
|
op.create_table(tmp_table_name, *new_columns)
|
|
meta.reflect()
|
|
|
|
try:
|
|
# copy data from the old table to the temp one
|
|
sql_select = sqlalchemy.sql.select(old_columns)
|
|
connection.execute(sqlalchemy.sql.insert(meta.tables[tmp_table_name])
|
|
.from_select(column_names, sql_select))
|
|
except Exception:
|
|
op.drop_table(tmp_table_name)
|
|
raise
|
|
|
|
# drop the old table and rename temp table to the old table name
|
|
op.drop_table(table_name)
|
|
op.rename_table(tmp_table_name, table_name)
|
|
|
|
# (re-)create indexes
|
|
for index in indexes:
|
|
op.create_index(op.f(index[0]), index[1], index[2], unique=index[3])
|