a54fb1b23a
This change ensures there is better messaging around the encrypted token migration, including a new phase to use for new installations, and fixes an issue encountered when running database migrations for new installations
140 lines
3.6 KiB
Python
140 lines
3.6 KiB
Python
import json
|
|
import logging
|
|
import uuid
|
|
|
|
from abc import ABCMeta, abstractmethod, abstractproperty
|
|
from datetime import datetime
|
|
from six import add_metaclass
|
|
|
|
from alembic import op
|
|
from sqlalchemy import text
|
|
|
|
from util.abchelpers import nooper
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def escape_table_name(table_name):
|
|
if op.get_bind().engine.name == 'postgresql':
|
|
# Needed for the `user` table.
|
|
return '"%s"' % table_name
|
|
|
|
return table_name
|
|
|
|
|
|
class DataTypes(object):
|
|
@staticmethod
|
|
def DateTime():
|
|
return datetime.now()
|
|
|
|
@staticmethod
|
|
def Date():
|
|
return datetime.now()
|
|
|
|
@staticmethod
|
|
def String():
|
|
return 'somestringvalue'
|
|
|
|
@staticmethod
|
|
def Token():
|
|
return '%s%s' % ('a' * 60, 'b' * 60)
|
|
|
|
@staticmethod
|
|
def UTF8Char():
|
|
return 'some other value'
|
|
|
|
@staticmethod
|
|
def UUID():
|
|
return str(uuid.uuid4())
|
|
|
|
@staticmethod
|
|
def JSON():
|
|
return json.dumps(dict(foo='bar', baz='meh'))
|
|
|
|
@staticmethod
|
|
def Boolean():
|
|
if op.get_bind().engine.name == 'postgresql':
|
|
return True
|
|
|
|
return 1
|
|
|
|
@staticmethod
|
|
def BigInteger():
|
|
return 21474836470
|
|
|
|
@staticmethod
|
|
def Integer():
|
|
return 42
|
|
|
|
@staticmethod
|
|
def Constant(value):
|
|
def get_value():
|
|
return value
|
|
return get_value
|
|
|
|
@staticmethod
|
|
def Foreign(table_name):
|
|
def get_index():
|
|
result = op.get_bind().execute("SELECT id FROM %s LIMIT 1" % escape_table_name(table_name))
|
|
try:
|
|
return list(result)[0][0]
|
|
except IndexError:
|
|
raise Exception('Could not find row for table %s' % table_name)
|
|
finally:
|
|
result.close()
|
|
|
|
return get_index
|
|
|
|
|
|
@add_metaclass(ABCMeta)
|
|
class MigrationTester(object):
|
|
""" Implements an interface for adding testing capabilities to the
|
|
data model migration system in Alembic.
|
|
"""
|
|
TestDataType = DataTypes
|
|
|
|
@abstractproperty
|
|
def is_testing(self):
|
|
""" Returns whether we are currently under a migration test. """
|
|
|
|
@abstractmethod
|
|
def populate_table(self, table_name, fields):
|
|
""" Called to populate a table with the given fields filled in with testing data. """
|
|
|
|
@abstractmethod
|
|
def populate_column(self, table_name, col_name, field_type):
|
|
""" Called to populate a column in a table to be filled in with testing data. """
|
|
|
|
|
|
@nooper
|
|
class NoopTester(MigrationTester):
|
|
""" No-op version of the tester, designed for production workloads. """
|
|
|
|
|
|
class PopulateTestDataTester(MigrationTester):
|
|
@property
|
|
def is_testing(self):
|
|
return True
|
|
|
|
def populate_table(self, table_name, fields):
|
|
columns = {field_name: field_type() for field_name, field_type in fields}
|
|
field_name_vars = [':' + field_name for field_name, _ in fields]
|
|
|
|
if op.get_bind().engine.name == 'postgresql':
|
|
field_names = ["%s" % field_name for field_name, _ in fields]
|
|
else:
|
|
field_names = ["`%s`" % field_name for field_name, _ in fields]
|
|
|
|
table_name = escape_table_name(table_name)
|
|
query = text('INSERT INTO %s (%s) VALUES (%s)' % (table_name, ', '.join(field_names),
|
|
', '.join(field_name_vars)))
|
|
logger.info("Executing test query %s with values %s", query, columns.values())
|
|
op.get_bind().execute(query, **columns)
|
|
|
|
def populate_column(self, table_name, col_name, field_type):
|
|
col_value = field_type()
|
|
row_id = DataTypes.Foreign(table_name)()
|
|
|
|
table_name = escape_table_name(table_name)
|
|
update_text = text("UPDATE %s SET %s=:col_value where ID=:row_id" % (table_name, col_name))
|
|
logger.info("Executing test query %s with value %s on row %s", update_text, col_value, row_id)
|
|
op.get_bind().execute(update_text, col_value=col_value, row_id=row_id)
|