Add charset/snapshot_file to vttest. (vitessio#2611)

This helps for testing efforts that require restoring from a snapshot. Also, setting charset is already allowed via start_vt_processes, so this PR just exposes it through run_local_database.

Author: thompsonja

java/client/src/test/java/com/youtube/vitess/client/TestUtil.java

@@ -47,10 +47,10 @@ public static void setupTestEnv(TestEnv testEnv) throws Exception {
         continue;
       }
       try {
-        Type mapType = new TypeToken<Map<String, Integer>>() {}.getType();
-        Map<String, Integer> map = new Gson().fromJson(line, mapType);
+        Type mapType = new TypeToken<Map<String, Object>>() {}.getType();
+        Map<String, Object> map = new Gson().fromJson(line, mapType);
         testEnv.setPythonScriptProcess(p);
-        testEnv.setPort(map.get(System.getProperty(PROPERTY_KEY_CLIENT_TEST_PORT)));
+        testEnv.setPort(((Double)map.get(System.getProperty(PROPERTY_KEY_CLIENT_TEST_PORT))).intValue());
         return;
       } catch (JsonSyntaxException e) {
         logger.error("JsonSyntaxException parsing setup command output: " + line, e);

py/vttest/local_database.py

@@ -21,7 +21,9 @@ def __init__(self,
                web_dir=None,
                default_schema_dir=None,
                extra_my_cnf=None,
-               web_dir2=None):
+               web_dir2=None,
+               snapshot_file=None,
+               charset='utf8'):
     """Initializes an object of this class.
 
     Args:
@@ -41,6 +43,8 @@ def __init__(self,
       extra_my_cnf: additional cnf file to use for the EXTRA_MY_CNF var.
       web_dir2: see the documentation for the corresponding command line
           flag in run_local_database.py
+      snapshot_file: A MySQL DB snapshot file.
+      charset: MySQL charset.
     """
 
     self.topology = topology
@@ -51,17 +55,20 @@ def __init__(self,
     self.default_schema_dir = default_schema_dir
     self.extra_my_cnf = extra_my_cnf
     self.web_dir2 = web_dir2
+    self.snapshot_file = snapshot_file
+    self.charset = charset
 
   def setup(self):
     """Create a MySQL instance and all Vitess processes."""
     mysql_port = environment.get_port('mysql')
     self.directory = environment.get_test_directory()
     self.mysql_db = environment.mysql_db_class(
-        self.directory, mysql_port, self.extra_my_cnf)
+        self.directory, mysql_port, self.extra_my_cnf, self.snapshot_file)
 
     self.mysql_db.setup()
-    self.create_databases()
-    self.load_schema()
+    if not self.snapshot_file:
+      self.create_databases()
+      self.load_schema()
     if self.init_data_options is not None:
       self.rng = random.Random(self.init_data_options.rng_seed)
       self.populate_with_random_data()
@@ -70,7 +77,7 @@ def setup(self):
 
     vt_processes.start_vt_processes(self.directory, self.topology,
                                     self.mysql_db, self.schema_dir,
-                                    web_dir=self.web_dir,
+                                    charset=self.charset, web_dir=self.web_dir,
                                     web_dir2=self.web_dir2)
 
   def teardown(self):
@@ -105,7 +112,9 @@ def config(self):
 
     result = {
         'port': vt_processes.vtcombo_process.port,
-        }
+        'socket': self.mysql_db.unix_socket(),
+    }
+
     if environment.get_protocol() == 'grpc':
       result['grpc_port'] = vt_processes.vtcombo_process.grpc_port
     return result

py/vttest/mysql_db.py

@@ -9,10 +9,11 @@
 class MySqlDB(object):
   """A MySqlDB contains basic info about a MySQL instance."""
 
-  def __init__(self, directory, port, extra_my_cnf=None):
+  def __init__(self, directory, port, extra_my_cnf=None, snapshot_file=None):
     self._directory = directory
     self._port = port
     self._extra_my_cnf = extra_my_cnf
+    self._snapshot_file = snapshot_file
 
   def setup(self, port):
     """Starts the MySQL database."""

py/vttest/mysql_db_mysqlctl.py

@@ -18,8 +18,9 @@
 class MySqlDBMysqlctl(mysql_db.MySqlDB):
   """Contains data and methods to manage a MySQL instance using mysqlctl."""
 
-  def __init__(self, directory, port, extra_my_cnf):
-    super(MySqlDBMysqlctl, self).__init__(directory, port, extra_my_cnf)
+  def __init__(self, directory, port, extra_my_cnf, snapshot_file=None):
+    super(MySqlDBMysqlctl, self).__init__(
+        directory, port, extra_my_cnf, snapshot_file)
 
   def setup(self):
     cmd = [

py/vttest/run_local_database.py

@@ -34,18 +34,17 @@
 import os
 import sys
 
-from vtdb import prefer_vtroot_imports  # pylint: disable=unused-import
 
 from google.protobuf import text_format
 
+from vtproto import vttest_pb2
+from vtdb import prefer_vtroot_imports  # pylint: disable=unused-import
 from vttest import environment
+from vttest import init_data_options
 from vttest import local_database
 from vttest import mysql_flavor
-from vttest import init_data_options
 from vttest import sharding_utils
 
-from vtproto import vttest_pb2
-
 
 def main(cmdline_options):
   topology = vttest_pb2.VTTestTopology()
@@ -96,7 +95,9 @@ def main(cmdline_options):
       web_dir=cmdline_options.web_dir,
       web_dir2=cmdline_options.web_dir2,
       default_schema_dir=cmdline_options.default_schema_dir,
-      extra_my_cnf=extra_my_cnf) as local_db:
+      extra_my_cnf=extra_my_cnf,
+      charset=cmdline_options.charset,
+      snapshot_file=cmdline_options.snapshot_file) as local_db:
     print json.dumps(local_db.config())
     sys.stdout.flush()
     try:
@@ -186,6 +187,9 @@ def main(cmdline_options):
                     help='Replica tablets per shard (includes master)')
   parser.add_option('--rdonly_count', type='int', default=1,
                     help='Rdonly tablets per shard')
+  parser.add_option('--charset', default='utf8', help='MySQL charset')
+  parser.add_option(
+      '--snapshot_file', default=None, help='A MySQL DB snapshot file')
   (options, args) = parser.parse_args()
   if options.verbose:
     logging.getLogger().setLevel(logging.DEBUG)