345 lines
13 KiB
Python
345 lines
13 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2022 Google LLC. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Command for spanner samples init."""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import unicode_literals
|
|
|
|
import json
|
|
import os
|
|
import textwrap
|
|
|
|
from apitools.base.py import exceptions as apitools_exceptions
|
|
from googlecloudsdk.api_lib.spanner import database_operations
|
|
from googlecloudsdk.api_lib.spanner import database_sessions
|
|
from googlecloudsdk.api_lib.spanner import databases
|
|
from googlecloudsdk.api_lib.spanner import instances
|
|
from googlecloudsdk.api_lib.storage import storage_api
|
|
from googlecloudsdk.api_lib.storage import storage_util
|
|
from googlecloudsdk.calliope import base
|
|
from googlecloudsdk.calliope import exceptions as calliope_exceptions
|
|
from googlecloudsdk.command_lib.spanner import ddl_parser
|
|
from googlecloudsdk.command_lib.spanner import samples
|
|
from googlecloudsdk.core import log
|
|
from googlecloudsdk.core import properties
|
|
from googlecloudsdk.core import resources
|
|
from googlecloudsdk.core.console import progress_tracker
|
|
from googlecloudsdk.core.util import files
|
|
from googlecloudsdk.core.util import retry
|
|
|
|
|
|
def check_instance(instance_id):
|
|
"""Raise if the given instance doesn't exist."""
|
|
try:
|
|
instances.Get(instance_id)
|
|
except apitools_exceptions.HttpNotFoundError:
|
|
raise ValueError(
|
|
textwrap.dedent("""\
|
|
Instance '{instance_id}' does not exist. Create it with:
|
|
|
|
$ gcloud spanner instances create {instance_id}
|
|
""".format(instance_id=instance_id)))
|
|
|
|
|
|
def download_sample_files(appname):
|
|
"""Download schema and binaries for the given sample app.
|
|
|
|
If the schema and all binaries exist already, skip download. If any file
|
|
doesn't exist, download them all.
|
|
|
|
Args:
|
|
appname: The name of the sample app, should exist in samples.APP_NAMES
|
|
"""
|
|
storage_client = storage_api.StorageClient()
|
|
bucket_ref = storage_util.BucketReference.FromUrl(samples.GCS_BUCKET)
|
|
|
|
# Get the GCS object ref and local path for each file
|
|
gcs_to_local = [(storage_util.ObjectReference.FromBucketRef(
|
|
bucket_ref, samples.get_gcs_schema_name(appname)),
|
|
samples.get_local_schema_path(appname))]
|
|
gcs_bin_msgs = storage_client.ListBucket(
|
|
bucket_ref, prefix=samples.get_gcs_bin_prefix(appname))
|
|
bin_path = samples.get_local_bin_path(appname)
|
|
for gcs_ref in gcs_bin_msgs:
|
|
# Skip folder or dir result in ListBucket result.
|
|
if not gcs_ref.name.split('/')[-1]:
|
|
continue
|
|
gcs_ref = storage_util.ObjectReference.FromMessage(gcs_ref)
|
|
local_path = os.path.join(bin_path, gcs_ref.name.split('/')[-1])
|
|
gcs_to_local.append((gcs_ref, local_path))
|
|
|
|
if samples.has_sample_data_statements(appname):
|
|
insert_path = samples.get_gcs_data_insert_statements_prefix(appname)
|
|
gcs_insert_files = storage_client.ListBucket(bucket_ref, prefix=insert_path)
|
|
for insert_file in gcs_insert_files:
|
|
insert_file_ref = storage_util.ObjectReference.FromMessage(insert_file)
|
|
# Skip folder or dir in ListBucket result. Cannot use `os.path.isdir` to
|
|
# check due to GCS file naming convention.
|
|
if insert_file_ref.name.endswith('/'):
|
|
continue
|
|
data_local_path = samples.get_local_data_insert_statements_path(appname)
|
|
local_path = os.path.join(
|
|
data_local_path, insert_file_ref.name.split('/')[-1]
|
|
)
|
|
gcs_to_local.append((insert_file_ref, local_path))
|
|
|
|
# Download all files again if any file is missing
|
|
if any(not os.path.exists(file_path) for _, file_path in gcs_to_local):
|
|
log.status.Print('Downloading files for the {} sample app'.format(appname))
|
|
for gcs_ref, local_path in gcs_to_local:
|
|
log.status.Print('Downloading {}'.format(local_path))
|
|
local_dir = os.path.split(local_path)[0]
|
|
if not os.path.exists(local_dir):
|
|
files.MakeDir(local_dir)
|
|
storage_client.CopyFileFromGCS(gcs_ref, local_path, overwrite=True)
|
|
|
|
|
|
def _create_db_op(instance_ref, database_id, statements, database_dialect):
|
|
"""Wrapper over databases.Create with error handling."""
|
|
try:
|
|
return databases.Create(
|
|
instance_ref,
|
|
database_id,
|
|
statements,
|
|
database_dialect=database_dialect)
|
|
except apitools_exceptions.HttpConflictError:
|
|
raise ValueError(
|
|
textwrap.dedent("""\
|
|
Database '{database_id}' exists already. Delete it with:
|
|
|
|
$ gcloud spanner databases delete {database_id} --instance={instance_id}
|
|
""".format(
|
|
database_id=database_id, instance_id=instance_ref.instancesId)))
|
|
except apitools_exceptions.HttpError as ex:
|
|
raise ValueError(json.loads(ex.content)['error']['message'])
|
|
except Exception: # pylint: disable=broad-except
|
|
raise ValueError("Failed to create database '{}'.".format(database_id))
|
|
|
|
|
|
def insert_sample_data_in_one_file(appname, file_name, session_ref):
|
|
"""Read and execute all insert statements in one file."""
|
|
if not samples.has_sample_data_statements(appname):
|
|
raise ValueError('{} cannot pre-populate data.'.format(appname))
|
|
|
|
insert_statements = files.ReadFileContents(file_name)
|
|
for insert_statement in insert_statements.split('\n'):
|
|
if not insert_statement:
|
|
continue
|
|
if not insert_statement.startswith('INSERT'):
|
|
continue
|
|
# Use a retryer to handle 409 txn abort errors that tend to happen
|
|
# when db is just created and group assignment contends with insert and
|
|
# commit dual-trip txns.
|
|
retry.Retryer(max_retrials=5).RetryOnException(
|
|
database_sessions.ExecuteSql,
|
|
args=[insert_statement, 'NORMAL', session_ref],
|
|
should_retry_if=lambda exc_type, *args: True,
|
|
sleep_ms=2000,
|
|
)
|
|
|
|
|
|
def insert_sample_data(appname, database_id, session_ref):
|
|
"""Insert sample data."""
|
|
if not samples.has_sample_data_statements(appname):
|
|
raise ValueError('{} cannot pre-populate data.'.format(appname))
|
|
|
|
with progress_tracker.ProgressTracker(
|
|
'Populating data into `{}`'.format(database_id),
|
|
aborted_message='Aborting wait for data population.\n',
|
|
):
|
|
data_files = files.GetDirectoryTreeListing(
|
|
samples.get_local_data_insert_statements_path(appname)
|
|
)
|
|
for data_file in data_files:
|
|
insert_sample_data_in_one_file(
|
|
appname,
|
|
data_file,
|
|
session_ref,
|
|
)
|
|
|
|
|
|
def check_create_db(appname, instance_ref, database_id):
|
|
"""Create the DB if it doesn't exist already, raise otherwise."""
|
|
schema_file = samples.get_local_schema_path(appname)
|
|
database_dialect = samples.get_database_dialect(appname)
|
|
|
|
schema = files.ReadFileContents(schema_file)
|
|
# Special case for POSTGRESQL dialect:
|
|
# a. CreateDatabase does not support additional_statements. Instead a
|
|
# separate call to UpdateDDL is used.
|
|
# b. ddl_parser only supports GSQL; instead remove comment lines, then
|
|
# split on ';'.
|
|
if database_dialect == databases.DATABASE_DIALECT_POSTGRESQL:
|
|
create_ddl = []
|
|
# Remove comments
|
|
schema = '\n'.join(
|
|
[line for line in schema.split('\n') if not line.startswith('--')])
|
|
# TODO(b/195711543): This would be incorrect if ';' is inside strings
|
|
# and / or comments.
|
|
update_ddl = [stmt for stmt in schema.split(';') if stmt]
|
|
else:
|
|
create_ddl = ddl_parser.PreprocessDDLWithParser(schema)
|
|
update_ddl = []
|
|
|
|
create_op = _create_db_op(instance_ref, database_id, create_ddl,
|
|
database_dialect)
|
|
database_operations.Await(create_op,
|
|
"Creating database '{}'".format(database_id))
|
|
if update_ddl:
|
|
database_ref = resources.REGISTRY.Parse(
|
|
database_id,
|
|
params={
|
|
'instancesId': instance_ref.instancesId,
|
|
'projectsId': instance_ref.projectsId,
|
|
},
|
|
collection='spanner.projects.instances.databases')
|
|
update_op = databases.UpdateDdl(database_ref, update_ddl)
|
|
database_operations.Await(update_op,
|
|
"Updating database '{}'".format(database_id))
|
|
|
|
|
|
@base.DefaultUniverseOnly
|
|
class Init(base.Command):
|
|
"""Initialize a Cloud Spanner sample app.
|
|
|
|
This command creates a Cloud Spanner database in the given instance for the
|
|
sample app and loads any initial data required by the application.
|
|
"""
|
|
|
|
detailed_help = {
|
|
'EXAMPLES': textwrap.dedent("""\
|
|
To initialize the 'finance' sample app using instance
|
|
'my-instance', run:
|
|
|
|
$ {command} finance --instance-id=my-instance
|
|
|
|
To initialize the 'finance-graph' sample app using instance
|
|
'my-instance', run:
|
|
|
|
$ {command} finance-graph --instance-id=my-instance
|
|
"""),
|
|
}
|
|
|
|
@staticmethod
|
|
def Args(parser):
|
|
"""Args is called by calliope to gather arguments for this command.
|
|
|
|
Args:
|
|
parser: An argparse parser that you can use to add arguments that go on
|
|
the command line after this command. Positional arguments are allowed.
|
|
"""
|
|
parser.add_argument(
|
|
'appname', help='The sample app name, e.g. "finance", "finance-graph".'
|
|
)
|
|
parser.add_argument(
|
|
'--instance-id',
|
|
required=True,
|
|
type=str,
|
|
help='The Cloud Spanner instance ID for the sample app.')
|
|
parser.add_argument(
|
|
'--database-id',
|
|
type=str,
|
|
help='ID of the new Cloud Spanner database to create for the sample '
|
|
'app.')
|
|
|
|
def Run(self, args):
|
|
"""This is what gets called when the user runs this command.
|
|
|
|
Args:
|
|
args: an argparse namespace. All the arguments that were provided to this
|
|
command invocation.
|
|
|
|
Returns:
|
|
Some value that we want to have printed later.
|
|
"""
|
|
appname = args.appname
|
|
try:
|
|
samples.check_appname(appname)
|
|
except ValueError as ex:
|
|
raise calliope_exceptions.BadArgumentException('APPNAME', ex)
|
|
|
|
instance_id = args.instance_id
|
|
instance_ref = resources.REGISTRY.Parse(
|
|
instance_id,
|
|
params={
|
|
'projectsId': properties.VALUES.core.project.GetOrFail,
|
|
},
|
|
collection='spanner.projects.instances')
|
|
|
|
if args.database_id is not None:
|
|
database_id = args.database_id
|
|
else:
|
|
database_id = samples.get_db_id_for_app(appname)
|
|
|
|
# Check that the instance exists
|
|
log.status.Print("Checking instance '{}'".format(instance_id))
|
|
try:
|
|
check_instance(instance_id)
|
|
except ValueError as ex:
|
|
raise calliope_exceptions.BadArgumentException('--instance-id', ex)
|
|
|
|
# Download any missing sample app binaries from GCS, including the schema
|
|
# file we need to create the DB
|
|
download_sample_files(appname)
|
|
|
|
# Create the sample app DB
|
|
log.status.Print(
|
|
"Initializing database '{database_id}' for sample app '{appname}'"
|
|
.format(database_id=database_id, appname=appname))
|
|
try:
|
|
check_create_db(appname, instance_ref, database_id)
|
|
except ValueError as ex:
|
|
raise calliope_exceptions.BadArgumentException('--database-id', ex)
|
|
|
|
if samples.has_sample_data_statements(appname):
|
|
database_ref = resources.REGISTRY.Parse(
|
|
database_id,
|
|
params={
|
|
'instancesId': instance_ref.instancesId,
|
|
'projectsId': instance_ref.projectsId,
|
|
},
|
|
collection='spanner.projects.instances.databases',
|
|
)
|
|
session = database_sessions.Create(database_ref)
|
|
session_ref = resources.REGISTRY.ParseRelativeName(
|
|
relative_name=session.name,
|
|
collection='spanner.projects.instances.databases.sessions',
|
|
)
|
|
try:
|
|
insert_sample_data(appname, database_id, session_ref)
|
|
except Exception:
|
|
raise SystemError(
|
|
'Failed to insert data for the database. Please fallback to '
|
|
'manually insert.'
|
|
)
|
|
else:
|
|
return textwrap.dedent("""\
|
|
Initialization done for your Spanner database.
|
|
""")
|
|
finally:
|
|
database_sessions.Delete(session_ref)
|
|
else:
|
|
backend_args = '{appname} --instance-id={instance_id}'.format(
|
|
appname=appname, instance_id=instance_id
|
|
)
|
|
if args.database_id is not None:
|
|
backend_args += ' --database-id {}'.format(database_id)
|
|
return textwrap.dedent("""\
|
|
Initialization done. Next, start the backend gRPC service with:
|
|
|
|
$ gcloud spanner samples backend {}
|
|
""".format(backend_args))
|