1import commands
2import os
3import sys
4
5
6GET_VERSION_QUERY = "SELECT latest_revision FROM application_meta LIMIT 1;"
7NULL_VERSION = "0000_00_00__00_00_00"
8
9
10class Commands(object):
11  UPGRADE = "upgrade"
12  DOWNGRADE = "downgrade"
13
14  @classmethod
15  def all(cls):
16    return (
17      cls.UPGRADE,
18      cls.DOWNGRADE,
19    )
20
21
22def run_command(command):
23  exit_code, output = commands.getstatusoutput(command)
24
25  if exit_code != 0:
26    raise ValueError("Error running command: %s" % command)
27
28  return output
29
30
31def run_query(query):
32  return run_command(
33    "make db_shell EXTRA=\"-e '{query}'\"".format(
34      query=query,
35    )
36  )
37
38
39def parse_version_query(mysql_output):
40  print "MYSQL OUTPUT: %s" % mysql_output
41  lines = (mysql_output.splitlines())
42  if len(lines) == 3:
43    # table not yet populated
44    return NULL_VERSION
45  return lines[3]
46
47
48def get_db_version():
49  return parse_version_query(
50    run_query(GET_VERSION_QUERY)
51  )
52
53
54def get_upgrade_files_gt(db_version, sql_scripts_dir):
55  return sorted(
56    filter(
57      lambda fname: "upgrade" in fname and fname[:len(db_version)] > db_version,
58      os.listdir(
59        sql_scripts_dir,
60      )
61    )
62  )
63
64
65def get_downgrade_file(db_version):
66  if db_version == NULL_VERSION:
67    return None
68  return "%s_downgrade.sql" % db_version
69
70
71def get_previous_db_version(current_db_version, sql_scripts_dir):
72  try:
73    return sorted(
74      filter(
75        lambda fname: "downgrade" in fname and fname[:len(current_db_version)] < current_db_version,
76        os.listdir(
77          sql_scripts_dir,
78        )
79      )
80    )[-1].replace("_downgrade.sql", "")
81  except IndexError:
82    return NULL_VERSION
83
84
85def run_sql_file(sql_scripts_dir, sql_fname):
86  run_command(
87    "make db_shell < %s" % os.path.join(sql_scripts_dir, sql_fname)
88  )
89
90
91def update_db_version(canonical_version):
92  run_query(
93    """
94      START TRANSACTION;
95      TRUNCATE TABLE application_meta;
96      INSERT INTO application_meta (latest_revision) VALUES(\\"{canonical_version}\\");
97      COMMIT;
98    """.format(
99      canonical_version=canonical_version,
100    ).replace("\n", "")
101  )
102
103
104def fname_to_canonical_version(fname):
105  return fname.replace(
106    "_upgrade.sql",
107    "",
108  ).replace(
109    "_downgrade.sql",
110    "",
111  )
112
113
114def upgrade_db(sql_scripts_dir):
115  print "Upgrading..."
116  upgrade_files = get_upgrade_files_gt(
117    get_db_version(),
118    sql_scripts_dir,
119  )
120  if not upgrade_files:
121    print "Database is upgraded to the latest revision"
122    return
123  for upgrade_file in upgrade_files:
124    run_sql_file(sql_scripts_dir, upgrade_file)
125    update_db_version(
126      fname_to_canonical_version(upgrade_file),
127    )
128  print "Database version is now set to %s" % get_db_version()
129
130
131def downgrade_db(sql_scripts_dir):
132  print "Downgrading..."
133  current_version = get_db_version()
134  downgrade_file = get_downgrade_file(current_version)
135  if downgrade_file is None:
136    print "Database is already in an empty state"
137    return
138  run_sql_file(sql_scripts_dir, downgrade_file)
139  update_db_version(
140    get_previous_db_version(
141      current_version,
142      sql_scripts_dir,
143    ),
144  )
145
146
147if __name__ == "__main__":
148  try:
149    {
150      Commands.UPGRADE: upgrade_db,
151      Commands.DOWNGRADE: downgrade_db,
152    }[sys.argv[1]](sys.argv[2])
153  except (IndexError, KeyError):
154    print "Usage: python tools/upgrade_db.py [%s|%s] {sql_script_dir}" % Commands.all()
155    sys.exit(1)
156