feat: Add new gcloud commands, API clients, and third-party libraries across various services.

gcloud auth application-default login/google-cloud-sdk/lib/surface/spanner/samples/__init__.py

@@ -0,0 +1,38 @@
+# -*- 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 group for spanner samples."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.calliope import base
+
+
+class Samples(base.Group):
+  """Cloud Spanner sample apps.
+
+  Each Cloud Spanner sample application includes a backend gRPC service
+  backed by a Cloud Spanner database and a workload script that generates
+  service traffic.
+
+  These sample apps are open source and available at
+  https://github.com/GoogleCloudPlatform/cloud-spanner-samples.
+
+  To see a list of available sample apps, run:
+
+      $ gcloud spanner samples list
+  """
+  pass

gcloud auth application-default login/google-cloud-sdk/lib/surface/spanner/samples/backend.py

@@ -0,0 +1,208 @@
+# -*- 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 backend."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import os
+import textwrap
+
+from apitools.base.py import exceptions as apitools_exceptions
+from googlecloudsdk.api_lib.spanner import databases
+from googlecloudsdk.calliope import arg_parsers
+from googlecloudsdk.calliope import base
+from googlecloudsdk.calliope import exceptions as calliope_exceptions
+from googlecloudsdk.command_lib.spanner import samples
+from googlecloudsdk.core import execution_utils
+from googlecloudsdk.core import properties
+from googlecloudsdk.core import resources
+from surface.spanner.samples import init as samples_init
+
+
+def _get_logfile_name(appname):
+  return '{}-backend.log'.format(appname)
+
+
+def _get_popen_jar(appname):
+  if appname not in samples.APPS:
+    raise ValueError("Unknown sample app '{}'".format(appname))
+  return os.path.join(
+      samples.get_local_bin_path(appname), samples.APPS[appname].backend_bin)
+
+
+# Note: Currently all apps supported use the same flag definitions.
+# If there is a need for different flags by app in the future this logic can
+# move to: third_party/py/googlecloudsdk/command_lib/spanner/samples.py
+def _get_popen_args(project, appname, instance_id, database_id=None, port=None):
+  """Get formatted args for server command."""
+  if database_id is None:
+    database_id = samples.get_db_id_for_app(appname)
+  flags = [
+      '--spanner_project_id={}'.format(project),
+      '--spanner_instance_id={}'.format(instance_id),
+      '--spanner_database_id={}'.format(database_id)
+  ]
+  if port is not None:
+    flags.append('--port={}'.format(port))
+  if samples.get_database_dialect(
+      appname) == databases.DATABASE_DIALECT_POSTGRESQL:
+    flags.append('--spanner_use_pg')
+  return flags
+
+
+def run_backend(project,
+                appname,
+                instance_id,
+                database_id=None,
+                port=None,
+                capture_logs=False):
+  """Run the backend service executable for the given sample app.
+
+  Args:
+    project: str, Name of the GCP project.
+    appname: str, Name of the sample app.
+    instance_id: str, Cloud Spanner instance ID.
+    database_id: str, Cloud Spanner database ID.
+    port: int, Port to run the service on.
+    capture_logs: bool, Whether to save logs to disk or print to stdout.
+
+  Returns:
+    subprocess.Popen or execution_utils.SubprocessTimeoutWrapper, The running
+      subprocess.
+  """
+  proc_args = ['java', '-jar']
+  proc_args.append(_get_popen_jar(appname))
+  proc_args.extend(
+      _get_popen_args(project, appname, instance_id, database_id, port))
+  capture_logs_fn = (
+      os.path.join(samples.SAMPLES_LOG_PATH, '{}-backend.log'.format(appname))
+      if capture_logs else None)
+  return samples.run_proc(proc_args, capture_logs_fn)
+
+
+class Backend(base.Command):
+  """Run the backend gRPC service for the given Cloud Spanner sample app.
+
+  This command starts the backend gRPC service for the given sample
+  application. Before starting the service, create the database and load any
+  initial data with:
+
+      $ {parent_command} init APPNAME --instance-id=INSTANCE_ID
+
+  After starting the service, generate traffic with:
+
+      $ {parent_command} workload APPNAME
+
+  To run all three steps together, use:
+
+      $ {parent_command} run APPNAME --instance-id=INSTANCE_ID
+  """
+
+  detailed_help = {
+      'EXAMPLES':
+          textwrap.dedent("""\
+          To run the backend gRPC service for the 'finance' sample app using
+          instance 'my-instance', run:
+
+          $ {command} finance --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".')
+    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='The Cloud Spanner database ID for the sample app.')
+    parser.add_argument(
+        '--duration',
+        default='1h',
+        type=arg_parsers.Duration(),
+        help=('Duration of time allowed to run before stopping the service.'))
+    parser.add_argument(
+        '--port', type=int, help=('Port on which to receive gRPC requests.'))
+
+  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)
+
+    project = properties.VALUES.core.project.GetOrFail()
+
+    instance_id = args.instance_id
+    try:
+      samples_init.check_instance(instance_id)
+    except ValueError as ex:
+      raise calliope_exceptions.BadArgumentException('--instance-id', ex)
+
+    if args.database_id is not None:
+      database_id = args.database_id
+    else:
+      database_id = samples.get_db_id_for_app(appname)
+    database_ref = resources.REGISTRY.Parse(
+        database_id,
+        params={
+            'projectsId': project,
+            'instancesId': instance_id
+        },
+        collection='spanner.projects.instances.databases')
+    try:
+      databases.Get(database_ref)
+    except apitools_exceptions.HttpNotFoundError as ex:
+      if args.database_id is not None:
+        raise calliope_exceptions.BadArgumentException('--database-id', ex)
+      else:
+        raise samples.SpannerSamplesError(
+            "Database {} doesn't exist. Did you run `gcloud spanner samples "
+            'init` first?'.format(database_id))
+
+    proc = run_backend(project, appname, instance_id, args.database_id,
+                       args.port)
+    try:
+      with execution_utils.RaisesKeyboardInterrupt():
+        proc.wait(args.duration)
+    except KeyboardInterrupt:
+      proc.terminate()
+      return 'Backend gRPC service killed'
+    except execution_utils.TIMEOUT_EXPIRED_ERR:
+      proc.terminate()
+      return 'Backend gRPC service timed out after {duration}s'.format(
+          duration=args.duration)
+    return

gcloud auth application-default login/google-cloud-sdk/lib/surface/spanner/samples/init.py

@@ -0,0 +1,344 @@
+# -*- 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))

gcloud auth application-default login/google-cloud-sdk/lib/surface/spanner/samples/list.py

@@ -0,0 +1,59 @@
+# -*- 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 list."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import textwrap
+
+from googlecloudsdk.calliope import base
+from googlecloudsdk.command_lib.spanner import samples
+
+
+class List(base.ListCommand):
+  """List available sample applications."""
+
+  detailed_help = {
+      'EXAMPLES':
+          textwrap.dedent("""\
+          To list available sample applications, run:
+
+          $ {command}
+        """),
+  }
+
+  @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.
+    """
+    pass
+
+  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.
+    """
+    return list(sorted(samples.APPS))

gcloud auth application-default login/google-cloud-sdk/lib/surface/spanner/samples/run.py

@@ -0,0 +1,210 @@
+# -*- 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 run."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import textwrap
+import time
+
+from apitools.base.py import exceptions as apitools_exceptions
+from googlecloudsdk.api_lib.spanner import databases
+from googlecloudsdk.calliope import arg_parsers
+from googlecloudsdk.calliope import base
+from googlecloudsdk.calliope import exceptions as calliope_exceptions
+from googlecloudsdk.command_lib.spanner import samples
+from googlecloudsdk.core import execution_utils
+from googlecloudsdk.core import log
+from googlecloudsdk.core import properties
+from googlecloudsdk.core import resources
+from surface.spanner.samples import backend as samples_backend
+from surface.spanner.samples import init as samples_init
+from surface.spanner.samples import workload as samples_workload
+
+
+class Run(base.Command):
+  """Run the given Cloud Spanner sample app.
+
+  Each Cloud Spanner sample application includes a backend gRPC service
+  backed by a Cloud Spanner database and a workload script that generates
+  service traffic. This command creates and initializes the Cloud Spanner
+  database and runs both the backend service and workload script.
+
+  These sample apps are open source and available at
+  https://github.com/GoogleCloudPlatform/cloud-spanner-samples.
+
+  To see a list of available sample apps, run:
+
+      $ {parent_command} list
+  """
+
+  detailed_help = {
+      'EXAMPLES':
+          textwrap.dedent("""\
+          To run the 'finance' sample app using instance 'my-instance', run:
+
+          $ {command} finance --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".')
+    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.')
+    parser.add_argument(
+        '--duration',
+        default='1h',
+        type=arg_parsers.Duration(),
+        help=('Duration of time allowed to run the sample app before stopping '
+              'the service.'))
+    parser.add_argument(
+        '--cleanup',
+        action='store_true',
+        default=True,
+        help=('Delete the instance after running the sample app.'))
+    parser.add_argument(
+        '--skip-init',
+        action='store_true',
+        default=False,
+        help=('Use an existing database instead of creating a new one.'))
+
+  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
+    project = properties.VALUES.core.project.GetOrFail()
+    instance_ref = resources.REGISTRY.Parse(
+        instance_id,
+        params={
+            'projectsId': project,
+        },
+        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)
+    duration = args.duration
+    skip_init = args.skip_init
+
+    try:
+      samples_init.check_instance(instance_id)
+    except ValueError as ex:
+      raise calliope_exceptions.BadArgumentException('--instance-id', ex)
+    log.status.Print(
+        "Initializing database '{database_id}' for sample app '{appname}'"
+        .format(database_id=database_id, appname=appname))
+    if skip_init:
+      database_ref = resources.REGISTRY.Parse(
+          database_id,
+          params={
+              'instancesId': instance_id,
+              'projectsId': project
+          },
+          collection='spanner.projects.instances.databases')
+      try:
+        databases.Get(database_ref)
+      # --skip-init assumes the database exists already, raise if it doesn't.
+      except apitools_exceptions.HttpNotFoundError:
+        bad_flag = ('--instance-id'
+                    if args.database_id is None else '--database-id')
+        raise calliope_exceptions.BadArgumentException(
+            bad_flag, "Database '{database_id}' does not exist in instance "
+            "'{instance_id}'. Re-run this command without `--skip-init` to "
+            'create it.'.format(
+                database_id=database_id, instance_id=instance_id))
+    else:
+      try:
+        # Download any missing sample files and create the DB.
+        if self.ReleaseTrack() == base.ReleaseTrack.ALPHA:
+          samples_init.download_sample_files(args.appname)
+        samples_init.check_create_db(args.appname, instance_ref, database_id)
+      except ValueError as ex:
+        raise calliope_exceptions.BadArgumentException('--database-id', ex)
+
+    be_proc = samples_backend.run_backend(project, appname, instance_id,
+                                          database_id)
+    try:
+      be_proc.wait(2)
+      return (
+          'The {} sample app backend gRPC service failed to start, is another '
+          'instance already running?'.format(appname))
+    except execution_utils.TIMEOUT_EXPIRED_ERR:
+      pass
+
+    now = int(time.time())
+    later = now + duration
+    wl_proc = samples_workload.run_workload(appname, capture_logs=True)
+    # Wait a second to let the workload print startup logs
+    time.sleep(1)
+    log.status.Print(
+        '\nGenerating workload for database, start timestamp: {now}, end '
+        'timestamp: {later}. Press ^C to stop.'.format(now=now, later=later))
+
+    try:
+      with execution_utils.RaisesKeyboardInterrupt():
+        wl_proc.wait(duration)
+    except KeyboardInterrupt:
+      wl_proc.terminate()
+      be_proc.terminate()
+      log.status.Print('Backend gRPC service and workload generator killed')
+    except execution_utils.TIMEOUT_EXPIRED_ERR:
+      wl_proc.terminate()
+      be_proc.terminate()
+      log.status.Print(
+          'Backend gRPC service and workload generator killed after {duration}s'
+          .format(duration=duration))
+
+    if args.cleanup:
+      log.status.Print("Deleting database '{}'".format(database_id))
+      database_ref = resources.REGISTRY.Parse(
+          database_id,
+          params={
+              'projectsId': properties.VALUES.core.project.GetOrFail,
+              'instancesId': instance_ref.instancesId
+          },
+          collection='spanner.projects.instances.databases')
+      databases.Delete(database_ref)
+
+    log.status.Print('Done')
+    return

gcloud auth application-default login/google-cloud-sdk/lib/surface/spanner/samples/workload.py

@@ -0,0 +1,122 @@
+# -*- 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 workload."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import os
+import textwrap
+
+from googlecloudsdk.calliope import arg_parsers
+from googlecloudsdk.calliope import base
+from googlecloudsdk.command_lib.spanner import samples
+from googlecloudsdk.core import execution_utils
+
+
+def _get_popen_jar(appname):
+  if appname not in samples.APPS:
+    raise ValueError("Unknown sample app '{}'".format(appname))
+  return os.path.join(
+      samples.get_local_bin_path(appname), samples.APPS[appname].workload_bin)
+
+
+def run_workload(appname, port=None, capture_logs=False):
+  """Run the workload generator executable for the given sample app.
+
+  Args:
+    appname: str, Name of the sample app.
+    port: int, Port to run the service on.
+    capture_logs: bool, Whether to save logs to disk or print to stdout.
+
+  Returns:
+    subprocess.Popen or execution_utils.SubprocessTimeoutWrapper, The running
+      subprocess.
+  """
+  proc_args = ['java', '-jar', _get_popen_jar(appname)]
+  if port is not None:
+    proc_args.append('--port={}'.format(port))
+  capture_logs_fn = (
+      os.path.join(samples.SAMPLES_LOG_PATH, '{}-workload.log'.format(appname))
+      if capture_logs else None)
+  return samples.run_proc(proc_args, capture_logs_fn)
+
+
+class Workload(base.Command):
+  """Generate gRPC traffic for a given sample app's backend service.
+
+  Before sending traffic to the backend service, create the database and
+  start the service with:
+
+      $ {parent_command} init APPNAME --instance-id=INSTANCE_ID
+      $ {parent_command} backend APPNAME --instance-id=INSTANCE_ID
+
+  To run all three steps together, use:
+
+      $ {parent_command} run APPNAME --instance-id=INSTANCE_ID
+  """
+
+  detailed_help = {
+      'EXAMPLES':
+          textwrap.dedent("""\
+          To generate traffic for the 'finance' sample app, run:
+
+          $ {command} finance
+        """),
+  }
+
+  @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".')
+
+    parser.add_argument(
+        '--duration',
+        default='1h',
+        type=arg_parsers.Duration(),
+        help=('Duration of time allowed to run before stopping the workload.'))
+    parser.add_argument(
+        '--port', type=int, help=('Port of the running backend service.'))
+    parser.add_argument(
+        '--target-qps', type=int, help=('Target requests per second.'))
+
+  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.
+    """
+    proc = run_workload(args.appname, args.port)
+    try:
+      with execution_utils.RaisesKeyboardInterrupt():
+        return proc.wait(args.duration)
+    except KeyboardInterrupt:
+      proc.terminate()
+      return 'Workload generator killed'
+    except execution_utils.TIMEOUT_EXPIRED_ERR:
+      proc.terminate()
+      return 'Workload generator killed after {duration}s'.format(
+          duration=args.duration)
+    return