vitess-gh/test/keyspace_test.py

#!/usr/bin/env python
#
# Copyright 2013, Google Inc. All rights reserved.
# Use of this source code is governed by a BSD-style license that can
# be found in the LICENSE file.

import unittest

from vtdb import vtgate_client

import environment
import tablet
import utils
from protocols_flavor import protocols_flavor

SHARDED_KEYSPACE = 'TEST_KEYSPACE_SHARDED'
UNSHARDED_KEYSPACE = 'TEST_KEYSPACE_UNSHARDED'

# shards for SHARDED_KEYSPACE
# range '' - 80
shard_0_master = tablet.Tablet()
shard_0_replica = tablet.Tablet()
# range 80 - ''
shard_1_master = tablet.Tablet()
shard_1_replica = tablet.Tablet()

# shard for UNSHARDED_KEYSPACE
unsharded_master = tablet.Tablet()
unsharded_replica = tablet.Tablet()

shard_names = ['-80', '80-']
shard_kid_map = {
    '-80': [527875958493693904, 626750931627689502,
            345387386794260318, 332484755310826578,
            1842642426274125671, 1326307661227634652,
            1761124146422844620, 1661669973250483744,
            3361397649937244239, 2444880764308344533],
    '80-': [9767889778372766922, 9742070682920810358,
            10296850775085416642, 9537430901666854108,
            10440455099304929791, 11454183276974683945,
            11185910247776122031, 10460396697869122981,
            13379616110062597001, 12826553979133932576],
}

create_vt_insert_test = '''create table vt_insert_test (
id bigint auto_increment,
msg varchar(64),
keyspace_id bigint(20) unsigned NOT NULL,
primary key (id)
) Engine=InnoDB'''


def setUpModule():
  try:
    environment.topo_server().setup()

    setup_procs = [
        shard_0_master.init_mysql(),
        shard_0_replica.init_mysql(),
        shard_1_master.init_mysql(),
        shard_1_replica.init_mysql(),
        unsharded_master.init_mysql(),
        unsharded_replica.init_mysql(),
        ]
    utils.wait_procs(setup_procs)
    setup_tablets()
  except:
    tearDownModule()
    raise


def tearDownModule():
  if utils.options.skip_teardown:
    return

  tablet.kill_tablets([shard_0_master, shard_0_replica,
                       shard_1_master, shard_1_replica])
  teardown_procs = [
      shard_0_master.teardown_mysql(),
      shard_0_replica.teardown_mysql(),
      shard_1_master.teardown_mysql(),
      shard_1_replica.teardown_mysql(),
      unsharded_master.teardown_mysql(),
      unsharded_replica.teardown_mysql(),
      ]
  utils.wait_procs(teardown_procs, raise_on_error=False)

  environment.topo_server().teardown()
  utils.kill_sub_processes()
  utils.remove_tmp_files()

  shard_0_master.remove_tree()
  shard_0_replica.remove_tree()
  shard_1_master.remove_tree()
  shard_1_replica.remove_tree()
  unsharded_master.remove_tree()
  unsharded_replica.remove_tree()


def setup_tablets():
  setup_sharded_keyspace()
  setup_unsharded_keyspace()
  utils.VtGate().start()


def setup_sharded_keyspace():
  utils.run_vtctl(['CreateKeyspace', SHARDED_KEYSPACE])
  utils.run_vtctl(['SetKeyspaceShardingInfo', '-force', SHARDED_KEYSPACE,
                   'keyspace_id', 'uint64'])
  shard_0_master.init_tablet('master', keyspace=SHARDED_KEYSPACE, shard='-80')
  shard_0_replica.init_tablet(
      'replica', keyspace=SHARDED_KEYSPACE, shard='-80')
  shard_1_master.init_tablet('master', keyspace=SHARDED_KEYSPACE, shard='80-')
  shard_1_replica.init_tablet(
      'replica', keyspace=SHARDED_KEYSPACE, shard='80-')

  utils.run_vtctl(['RebuildKeyspaceGraph', SHARDED_KEYSPACE,], auto_log=True)

  for t in [shard_0_master, shard_0_replica, shard_1_master, shard_1_replica]:
    t.create_db('vt_test_keyspace_sharded')
    t.mquery(shard_0_master.dbname, create_vt_insert_test)
    t.start_vttablet(wait_for_state=None)

  for t in [shard_0_master, shard_0_replica, shard_1_master, shard_1_replica]:
    t.wait_for_vttablet_state('SERVING')

  utils.run_vtctl(['InitShardMaster', '%s/-80' % SHARDED_KEYSPACE,
                   shard_0_master.tablet_alias], auto_log=True)
  utils.run_vtctl(['InitShardMaster', '%s/80-' % SHARDED_KEYSPACE,
                   shard_1_master.tablet_alias], auto_log=True)

  utils.run_vtctl(['RebuildKeyspaceGraph', SHARDED_KEYSPACE],
                  auto_log=True)

  utils.check_srv_keyspace('test_nj', SHARDED_KEYSPACE,
                           'Partitions(master): -80 80-\n'
                           'Partitions(rdonly): -80 80-\n'
                           'Partitions(replica): -80 80-\n')


def setup_unsharded_keyspace():
  utils.run_vtctl(['CreateKeyspace', UNSHARDED_KEYSPACE])
  utils.run_vtctl(['SetKeyspaceShardingInfo', '-force', UNSHARDED_KEYSPACE,
                   'keyspace_id', 'uint64'])
  unsharded_master.init_tablet(
      'master', keyspace=UNSHARDED_KEYSPACE, shard='0')
  unsharded_replica.init_tablet(
      'replica', keyspace=UNSHARDED_KEYSPACE, shard='0')

  utils.run_vtctl(['RebuildKeyspaceGraph', UNSHARDED_KEYSPACE,], auto_log=True)

  for t in [unsharded_master, unsharded_replica]:
    t.create_db('vt_test_keyspace_unsharded')
    t.mquery(unsharded_master.dbname, create_vt_insert_test)
    t.start_vttablet(wait_for_state=None)

  for t in [unsharded_master, unsharded_replica]:
    t.wait_for_vttablet_state('SERVING')

  utils.run_vtctl(['InitShardMaster', '%s/0' % UNSHARDED_KEYSPACE,
                   unsharded_master.tablet_alias], auto_log=True)

  utils.run_vtctl(['RebuildKeyspaceGraph', UNSHARDED_KEYSPACE],
                  auto_log=True)

  utils.check_srv_keyspace('test_nj', UNSHARDED_KEYSPACE,
                           'Partitions(master): -\n'
                           'Partitions(rdonly): -\n'
                           'Partitions(replica): -\n')


ALL_DB_TYPES = ['master', 'rdonly', 'replica']


class TestKeyspace(unittest.TestCase):

  def _read_srv_keyspace(self, keyspace_name):
    addr = utils.vtgate.rpc_endpoint()
    protocol = protocols_flavor().vtgate_python_protocol()
    conn = vtgate_client.connect(protocol, addr, 30.0)
    result = conn.get_srv_keyspace(keyspace_name)
    conn.close()
    return result

  def test_get_keyspace(self):
    ki = utils.run_vtctl_json(['GetKeyspace', UNSHARDED_KEYSPACE])
    self.assertEqual('keyspace_id', ki['sharding_column_name'])
    self.assertEqual(1, ki['sharding_column_type'])

  def test_delete_keyspace(self):
    utils.run_vtctl(['CreateKeyspace', 'test_delete_keyspace'])
    utils.run_vtctl(['CreateShard', 'test_delete_keyspace/0'])
    utils.run_vtctl(
        ['InitTablet', '-keyspace=test_delete_keyspace', '-shard=0',
         'test_nj-0000000100', 'master'])

    # Can't delete keyspace if there are shards present.
    utils.run_vtctl(
        ['DeleteKeyspace', 'test_delete_keyspace'], expect_fail=True)
    # Can't delete shard if there are tablets present.
    utils.run_vtctl(['DeleteShard', 'test_delete_keyspace/0'], expect_fail=True)

    # Use recursive DeleteShard to remove tablets.
    utils.run_vtctl(['DeleteShard', '-recursive', 'test_delete_keyspace/0'])
    # Now non-recursive DeleteKeyspace should work.
    utils.run_vtctl(['DeleteKeyspace', 'test_delete_keyspace'])

    # Start over and this time use recursive DeleteKeyspace to do everything.
    utils.run_vtctl(['CreateKeyspace', 'test_delete_keyspace'])
    utils.run_vtctl(['CreateShard', 'test_delete_keyspace/0'])
    utils.run_vtctl(
        ['InitTablet', '-port=1234', '-keyspace=test_delete_keyspace',
         '-shard=0', 'test_nj-0000000100', 'master'])

    # Create the serving/replication entries and check that they exist,
    # so we can later check they're deleted.
    utils.run_vtctl(['RebuildKeyspaceGraph', 'test_delete_keyspace'])
    utils.run_vtctl(['RebuildShardGraph', 'test_delete_keyspace/0'])
    utils.run_vtctl(
        ['GetShardReplication', 'test_nj', 'test_delete_keyspace/0'])
    utils.run_vtctl(['GetSrvKeyspace', 'test_nj', 'test_delete_keyspace'])
    utils.run_vtctl(['GetSrvShard', 'test_nj', 'test_delete_keyspace/0'])
    utils.run_vtctl(
        ['GetEndPoints', 'test_nj', 'test_delete_keyspace/0', 'master'])

    # Recursive DeleteKeyspace
    utils.run_vtctl(['DeleteKeyspace', '-recursive', 'test_delete_keyspace'])

    # Check that everything is gone.
    utils.run_vtctl(['GetKeyspace', 'test_delete_keyspace'], expect_fail=True)
    utils.run_vtctl(['GetShard', 'test_delete_keyspace/0'], expect_fail=True)
    utils.run_vtctl(['GetTablet', 'test_nj-0000000100'], expect_fail=True)
    utils.run_vtctl(
        ['GetShardReplication', 'test_nj', 'test_delete_keyspace/0'],
        expect_fail=True)
    utils.run_vtctl(
        ['GetSrvKeyspace', 'test_nj', 'test_delete_keyspace'],
        expect_fail=True)
    utils.run_vtctl(
        ['GetSrvShard', 'test_nj', 'test_delete_keyspace/0'], expect_fail=True)
    utils.run_vtctl(
        ['GetEndPoints', 'test_nj', 'test_delete_keyspace/0', 'master'],
        expect_fail=True)

  def test_remove_keyspace_cell(self):
    utils.run_vtctl(['CreateKeyspace', 'test_delete_keyspace'])
    utils.run_vtctl(['CreateShard', 'test_delete_keyspace/0'])
    utils.run_vtctl(['CreateShard', 'test_delete_keyspace/1'])
    utils.run_vtctl(
        ['InitTablet', '-port=1234', '-keyspace=test_delete_keyspace',
         '-shard=0', 'test_ca-0000000100', 'master'])
    utils.run_vtctl(
        ['InitTablet', '-port=1234', '-keyspace=test_delete_keyspace',
         '-shard=0', 'test_nj-0000000100', 'replica'])
    utils.run_vtctl(
        ['InitTablet', '-port=1234', '-keyspace=test_delete_keyspace',
         '-shard=1', 'test_nj-0000000101', 'replica'])

    # Create the serving/replication entries and check that they exist,
    # so we can later check they're deleted.
    utils.run_vtctl(['RebuildKeyspaceGraph', 'test_delete_keyspace'])
    utils.run_vtctl(['RebuildShardGraph', 'test_delete_keyspace/0'])
    utils.run_vtctl(['RebuildShardGraph', 'test_delete_keyspace/1'])
    utils.run_vtctl(
        ['GetShardReplication', 'test_nj', 'test_delete_keyspace/0'])
    utils.run_vtctl(
        ['GetShardReplication', 'test_nj', 'test_delete_keyspace/1'])
    utils.run_vtctl(['GetSrvKeyspace', 'test_nj', 'test_delete_keyspace'])
    utils.run_vtctl(['GetSrvShard', 'test_nj', 'test_delete_keyspace/0'])
    utils.run_vtctl(['GetSrvShard', 'test_nj', 'test_delete_keyspace/1'])
    utils.run_vtctl(
        ['GetEndPoints', 'test_nj', 'test_delete_keyspace/0', 'replica'])
    utils.run_vtctl(
        ['GetEndPoints', 'test_nj', 'test_delete_keyspace/1', 'replica'])

    # Just remove the shard from one cell (including tablets),
    # but leaving the global records and other cells/shards alone.
    utils.run_vtctl(
        ['RemoveShardCell', '-recursive', 'test_delete_keyspace/0', 'test_nj'])
    utils.run_vtctl(['RebuildKeyspaceGraph', 'test_delete_keyspace'])
    utils.run_vtctl(['RebuildShardGraph', 'test_delete_keyspace/0'])

    utils.run_vtctl(['GetKeyspace', 'test_delete_keyspace'])
    utils.run_vtctl(['GetShard', 'test_delete_keyspace/0'])
    utils.run_vtctl(['GetTablet', 'test_ca-0000000100'])
    utils.run_vtctl(['GetTablet', 'test_nj-0000000100'], expect_fail=True)
    utils.run_vtctl(['GetTablet', 'test_nj-0000000101'])
    utils.run_vtctl(
        ['GetShardReplication', 'test_ca', 'test_delete_keyspace/0'])
    utils.run_vtctl(
        ['GetShardReplication', 'test_nj', 'test_delete_keyspace/0'],
        expect_fail=True)
    utils.run_vtctl(
        ['GetShardReplication', 'test_nj', 'test_delete_keyspace/1'])
    utils.run_vtctl(['GetSrvKeyspace', 'test_nj', 'test_delete_keyspace'])
    utils.run_vtctl(['GetSrvShard', 'test_nj', 'test_delete_keyspace/0'])
    utils.run_vtctl(
        ['GetEndPoints', 'test_nj', 'test_delete_keyspace/1', 'replica'])
    utils.run_vtctl(
        ['GetEndPoints', 'test_nj', 'test_delete_keyspace/0', 'replica'],
        expect_fail=True)

    # Add it back to do another test.
    utils.run_vtctl(
        ['InitTablet', '-port=1234', '-keyspace=test_delete_keyspace',
         '-shard=0', 'test_nj-0000000100', 'replica'])
    utils.run_vtctl(['RebuildKeyspaceGraph', 'test_delete_keyspace'])
    utils.run_vtctl(['RebuildShardGraph', 'test_delete_keyspace/0'])
    utils.run_vtctl(
        ['GetShardReplication', 'test_nj', 'test_delete_keyspace/0'])

    # Now use RemoveKeyspaceCell to remove all shards.
    utils.run_vtctl(
        ['RemoveKeyspaceCell', '-recursive', 'test_delete_keyspace',
         'test_nj'])
    utils.run_vtctl(['RebuildKeyspaceGraph', 'test_delete_keyspace'])
    utils.run_vtctl(['RebuildShardGraph', 'test_delete_keyspace/0'])
    utils.run_vtctl(['RebuildShardGraph', 'test_delete_keyspace/1'])

    utils.run_vtctl(
        ['GetShardReplication', 'test_ca', 'test_delete_keyspace/0'])
    utils.run_vtctl(
        ['GetShardReplication', 'test_nj', 'test_delete_keyspace/0'],
        expect_fail=True)
    utils.run_vtctl(
        ['GetShardReplication', 'test_nj', 'test_delete_keyspace/1'],
        expect_fail=True)

    # Clean up.
    utils.run_vtctl(['DeleteKeyspace', '-recursive', 'test_delete_keyspace'])

  def test_shard_count(self):
    sharded_ks = self._read_srv_keyspace(SHARDED_KEYSPACE)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(sharded_ks.get_shard_count(db_type), 2)
    unsharded_ks = self._read_srv_keyspace(UNSHARDED_KEYSPACE)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(unsharded_ks.get_shard_count(db_type), 1)

  def test_shard_names(self):
    sharded_ks = self._read_srv_keyspace(SHARDED_KEYSPACE)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(sharded_ks.get_shard_names(db_type), ['-80', '80-'])
    unsharded_ks = self._read_srv_keyspace(UNSHARDED_KEYSPACE)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(unsharded_ks.get_shard_names(db_type), ['0'])

  def test_keyspace_id_to_shard_name(self):
    sharded_ks = self._read_srv_keyspace(SHARDED_KEYSPACE)
    for _, sn in enumerate(shard_names):
      for keyspace_id in shard_kid_map[sn]:
        self.assertEqual(
            sharded_ks.keyspace_id_to_shard_name_for_db_type(keyspace_id,
                                                             'master'), sn)
    unsharded_ks = self._read_srv_keyspace(UNSHARDED_KEYSPACE)
    for keyspace_id in shard_kid_map[sn]:
      self.assertEqual(
          unsharded_ks.keyspace_id_to_shard_name_for_db_type(
              keyspace_id, 'master'),
          '0')

  def test_get_srv_keyspace_names(self):
    stdout, _ = utils.run_vtctl(['GetSrvKeyspaceNames', 'test_nj'],
                                trap_output=True)
    self.assertEqual(
        set(stdout.splitlines()), {SHARDED_KEYSPACE, UNSHARDED_KEYSPACE})


if __name__ == '__main__':
  utils.main()