1from autotest_lib.client.common_lib import global_config, host_protections
2
3def migrate_up(manager):
4    manager.execute_script(ADD_PROTECTION_COLUMN)
5
6def migrate_down(manager):
7    manager.execute(DROP_COLUMN)
8
9default_protection = global_config.global_config.get_config_value(
10    'HOSTS', 'default_protection')
11default_protection_value = host_protections.Protection.get_value(
12    default_protection)
13
14ADD_PROTECTION_COLUMN = """ALTER TABLE hosts
15                           ADD COLUMN protection INT NOT NULL
16                           DEFAULT %s;
17
18                           ALTER TABLE hosts
19                           ALTER COLUMN protection
20                           DROP DEFAULT;
21                           """ % default_protection_value
22
23DROP_COLUMN = """ALTER TABLE hosts
24                 DROP COLUMN protection"""
25