def create(self, **kwargs):
self.create_gke_cluster = (kwargs.get('create_gke_cluster',
'false').lower() != 'false')
if self.create_gke_cluster and 'GKE_NUM_NODES' not in kwargs:
raise base_environment.VitessEnvironmentError(
'Must specify GKE_NUM_NODES')
if 'GKE_CLUSTER_NAME' not in kwargs:
kwargs['GKE_CLUSTER_NAME'] = getpass.getuser()
if 'VITESS_NAME' not in kwargs:
kwargs['VITESS_NAME'] = getpass.getuser()
kwargs['TEST_MODE'] = '1'
self.script_dir = os.path.join(os.environ['VTTOP'],
'examples/kubernetes')
try:
subprocess.check_output(['gcloud', 'config', 'list'])
except OSError:
raise base_environment.VitessEnvironmentError(
'gcloud not found, please install by visiting cloud.google.com'
)
if 'project' in kwargs:
logging.info('Setting project to %s', kwargs['project'])
subprocess.check_output(
['gcloud', 'config', 'set', 'project', kwargs['project']])
project_name_json = json.loads(
subprocess.check_output(
['gcloud', 'config', 'list', 'project', '--format', 'json']))
project_name = project_name_json['core']['project']
logging.info('Current project name: %s', project_name)
for k, v in kwargs.iteritems():
os.environ[k] = v
if self.create_gke_cluster:
cluster_up_txt = subprocess.check_output(
[os.path.join(self.script_dir, 'cluster-up.sh')],
cwd=self.script_dir,
stderr=subprocess.STDOUT)
logging.info(cluster_up_txt)
vitess_up_output = subprocess.check_output(
[os.path.join(self.script_dir, 'vitess-up.sh')],
cwd=self.script_dir,
stderr=subprocess.STDOUT)
logging.info(vitess_up_output)
self.use_named(kwargs['VITESS_NAME'])
def poll_for_varz(self,
tablet_name,
varz,
timeout=60.0,
condition_fn=None,
converter=str,
condition_msg=None):
"""Polls for varz to exist, or match specific conditions, within a timeout.
Args:
tablet_name: the name of the process that we're trying to poll vars from.
varz: name of the vars to fetch from varz
timeout: number of seconds that we should attempt to poll for.
condition_fn: a function that takes the var as input, and returns a truthy
value if it matches the success conditions.
converter: function to convert varz value
condition_msg: string describing the conditions that we're polling for,
used for error messaging.
Raises:
VitessEnvironmentError: Raised if the varz conditions aren't met within
the given timeout.
Returns:
dict of requested varz.
"""
start_time = time.time()
while True:
if (time.time() - start_time) >= timeout:
timeout_error_msg = 'Timed out polling for varz.'
if condition_fn and condition_msg:
timeout_error_msg += ' Condition "%s" not met.' % condition_msg
raise base_environment.VitessEnvironmentError(
timeout_error_msg)
hostname = self.get_tablet_ip_port(tablet_name)
host_varz = subprocess.check_output([
'kubectl', 'exec', '-ti',
self.get_tablet_pod_name(tablet_name),
'--namespace=%s' % self.cluster_name, 'curl',
'%s/debug/vars' % hostname
])
if not host_varz:
continue
host_varz = json.loads(host_varz)
if condition_fn is None or condition_fn(host_varz):
return host_varz
def use_named(self, instance_name):
# Check to make sure kubectl exists
try:
subprocess.check_output(['kubectl'])
except OSError:
raise base_environment.VitessEnvironmentError(
'kubectl not found, please install by visiting kubernetes.io or '
'running gcloud components update kubectl if using compute engine.'
)
get_address_template = (
'{{if ge (len .status.loadBalancer) 1}}'
'{{index (index .status.loadBalancer.ingress 0) "ip"}}'
'{{end}}')
get_address_params = [
'kubectl', 'get', '-o', 'template', '--template',
get_address_template, 'service', '--namespace', instance_name
]
start_time = time.time()
vtctld_addr = ''
while time.time() - start_time < 60 and not vtctld_addr:
vtctld_addr = subprocess.check_output(get_address_params +
['vtctld'],
stderr=subprocess.STDOUT)
self.vtctl_addr = '%s:15999' % vtctld_addr
self.vtctl_helper = vtctl_helper.VtctlHelper('grpc', self.vtctl_addr)
self.cluster_name = instance_name
keyspaces = self.vtctl_helper.execute_vtctl_command(['GetKeyspaces'])
self.mobs = filter(None, keyspaces.split('\n'))
self.keyspaces = self.mobs
if not self.keyspaces:
raise base_environment.VitessEnvironmentError(
'Invalid environment, no keyspaces found')
self.num_shards = []
for keyspace in self.keyspaces:
keyspace_info = json.loads(
self.vtctl_helper.execute_vtctl_command(
['GetKeyspace', keyspace]))
if not keyspace_info:
self.num_shards.append(1)
else:
self.num_shards.append(keyspace_info['split_shard_count'])
# This assumes that all keyspaces use the same set of cells
self.cells = json.loads(
self.vtctl_helper.execute_vtctl_command([
'GetShard',
'%s/%s' % (self.keyspaces[0],
utils.get_shard_name(0, self.num_shards[0]))
]))['cells']
self.primary_cells = self.cells
self.replica_instances = []
self.rdonly_instances = []
# This assumes that all cells are equivalent for k8s environments.
all_tablets_in_a_cell = self.vtctl_helper.execute_vtctl_command(
['ListAllTablets', self.cells[0]])
all_tablets_in_a_cell = [
x.split(' ')
for x in filter(None, all_tablets_in_a_cell.split('\n'))
]
for index, keyspace in enumerate(self.keyspaces):
keyspace_tablets_in_cell = [
tablet for tablet in all_tablets_in_a_cell
if tablet[1] == keyspace
]
replica_tablets_in_cell = [
tablet for tablet in keyspace_tablets_in_cell
if tablet[3] == 'master' or tablet[3] == 'replica'
]
replica_instances = len(
replica_tablets_in_cell) / self.num_shards[index]
self.replica_instances.append(replica_instances)
self.rdonly_instances.append((len(keyspace_tablets_in_cell) /
self.num_shards[index]) -
replica_instances)
# Converts keyspace name and alias to number of instances
self.keyspace_alias_to_num_instances_dict = {}
for index, keyspace in enumerate(self.keyspaces):
self.keyspace_alias_to_num_instances_dict[keyspace] = {
'replica': int(self.replica_instances[index]),
'rdonly': int(self.rdonly_instances[index])
}
start_time = time.time()
self.vtgate_addrs = {}
self.vtgate_conns = {}
for cell in self.cells:
self.vtgate_addr = ''
while time.time() - start_time < 60 and not self.vtgate_addr:
vtgate_addr = subprocess.check_output(get_address_params +
['vtgate-%s' % cell],
stderr=subprocess.STDOUT)
self.vtgate_addrs[cell] = '%s:15001' % vtgate_addr
self.vtgate_conns[cell] = vtgate_client.connect(
protocols_flavor.protocols_flavor().vtgate_python_protocol(),
self.vtgate_addrs[cell], 60)
def use_named(self, instance_name):
# Check to make sure kubectl exists
try:
subprocess.check_output(['kubectl'])
except OSError:
raise base_environment.VitessEnvironmentError(
'kubectl not found, please install by visiting kubernetes.io or '
'running gcloud components update kubectl if using compute engine.'
)
vtctld_ip = kubernetes_components.get_forwarded_ip(
'vtctld', instance_name)
self.vtctl_addr = '%s:15999' % vtctld_ip
self.vtctl_helper = vtctl_helper.VtctlHelper('grpc', self.vtctl_addr)
self.cluster_name = instance_name
keyspaces = self.vtctl_helper.execute_vtctl_command(['GetKeyspaces'])
self.mobs = filter(None, keyspaces.split('\n'))
self.keyspaces = self.mobs
if not self.keyspaces:
raise base_environment.VitessEnvironmentError(
'Invalid environment, no keyspaces found')
self.num_shards = []
self.shards = []
for keyspace in self.keyspaces:
shards = json.loads(
self.vtctl_helper.execute_vtctl_command(
['FindAllShardsInKeyspace', keyspace]))
self.shards.append(shards)
self.num_shards.append(len(shards))
# This assumes that all keyspaces use the same set of cells
self.cells = json.loads(
self.vtctl_helper.execute_vtctl_command([
'GetShard',
'%s/%s' % (self.keyspaces[0], self.shards[0].keys()[0])
]))['cells']
self.primary_cells = self.cells
self.replica_instances = []
self.rdonly_instances = []
# This assumes that all cells are equivalent for k8s environments.
all_tablets_in_a_cell = self.vtctl_helper.execute_vtctl_command(
['ListAllTablets', self.cells[0]])
all_tablets_in_a_cell = [
x.split(' ')
for x in filter(None, all_tablets_in_a_cell.split('\n'))
]
for index, keyspace in enumerate(self.keyspaces):
keyspace_tablets_in_cell = [
tablet for tablet in all_tablets_in_a_cell
if tablet[1] == keyspace
]
replica_tablets_in_cell = [
tablet for tablet in keyspace_tablets_in_cell
if tablet[3] == 'master' or tablet[3] == 'replica'
]
replica_instances = len(
replica_tablets_in_cell) / self.num_shards[index]
self.replica_instances.append(replica_instances)
self.rdonly_instances.append((len(keyspace_tablets_in_cell) /
self.num_shards[index]) -
replica_instances)
# Converts keyspace name and alias to number of instances
self.keyspace_alias_to_num_instances_dict = {}
for index, keyspace in enumerate(self.keyspaces):
self.keyspace_alias_to_num_instances_dict[keyspace] = {
'replica': int(self.replica_instances[index]),
'rdonly': int(self.rdonly_instances[index])
}
self.vtgate_addrs = {}
for cell in self.cells:
vtgate_ip = kubernetes_components.get_forwarded_ip(
'vtgate-%s' % cell, instance_name)
self.vtgate_addrs[cell] = '%s:15991' % vtgate_ip
super(K8sEnvironment, self).use_named(instance_name)
