cidb_admin: bail out if there are uncommitted or local-only changes

BUG=chromium:441892
TEST=local test

Change-Id: Ic4e003127cc23fee7e96fcf35346eb5c55db41ab
Reviewed-on: https://chromium-review.googlesource.com/316491
Commit-Ready: Aviv Keshet <akeshet@chromium.org>
Tested-by: Aviv Keshet <akeshet@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
diff --git a/scripts/cidb_admin.py b/scripts/cidb_admin.py
index d4f592a..74bc021 100644
--- a/scripts/cidb_admin.py
+++ b/scripts/cidb_admin.py
@@ -12,6 +12,7 @@
 from chromite.lib import commandline
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_logging as logging
+from chromite.lib import git
 
 
 MIGRATE = 'migrate'
@@ -44,6 +45,20 @@
 
   logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
 
+  # Check that we have no uncommitted files, and that our checkout's HEAD is
+  # contained in a remote branch. This is to ensure that we don't accidentally
+  # run uncommitted migrations.
+  uncommitted_files = git.RunGit(os.getcwd(), ['status', '-s']).output
+  if uncommitted_files:
+    cros_build_lib.Die('You appear to have uncommitted files. Aborting!')
+
+  remote_branches = git.RunGit(
+      os.getcwd(), ['branch', '-r', '--contains']).output
+  if not remote_branches:
+    cros_build_lib.Die(
+        'You appear to be on a local branch of chromite. Aborting!')
+
+
   if options.command == MIGRATE:
     positive_confirmation = 'please modify my database'
     warn = ('This option will apply schema changes to your existing database. '
@@ -61,14 +76,12 @@
                 os.path.join(options.cred_dir, 'host.txt'),
                 positive_confirmation)
   else:
-    print('No command or unsupported command. Exiting.')
-    exit()
+    cros_build_lib.Die('No command or unsupported command. Exiting.')
 
   print(warn)
   conf_string = cros_build_lib.GetInput('(%s)?: ' % positive_confirmation)
   if conf_string != positive_confirmation:
-    print('You changed your mind. Aborting.')
-    exit()
+    cros_build_lib.Die('You changed your mind. Aborting.')
 
   if options.command == MIGRATE:
     print('OK, applying migrations...')