# -*- coding: utf-8 -*- # # Copyright 2016 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 databases query.""" from __future__ import absolute_import from __future__ import division from __future__ import unicode_literals from googlecloudsdk.api_lib.spanner import database_sessions from googlecloudsdk.api_lib.util import apis from googlecloudsdk.calliope import arg_parsers from googlecloudsdk.calliope import base from googlecloudsdk.calliope import exceptions as c_exceptions from googlecloudsdk.command_lib.spanner import resource_args from googlecloudsdk.command_lib.spanner import sql from googlecloudsdk.command_lib.spanner.sql import QueryHasDml from googlecloudsdk.command_lib.util.apis import arg_utils from googlecloudsdk.core import log from googlecloudsdk.core import resources DETAILED_HELP = { 'EXAMPLES': """\ To execute a SQL SELECT statement against example-database under example-instance, run: $ {command} example-database --instance=example-instance --sql='SELECT * FROM MyTable WHERE MyKey = 1' """, } def CreateSession(args, database_role=None): """Creates a session. Args: args: an argparse namespace. All the arguments that were provided to the command invocation. database_role: Cloud Spanner database role which owns this session. Returns: A session reference to be used to execute the sql. """ session_name = database_sessions.Create(args.CONCEPTS.database.Parse(), database_role) return resources.REGISTRY.ParseRelativeName( relative_name=session_name.name, collection='spanner.projects.instances.databases.sessions') def AddBaseArgs(parser): """Parses provided arguments to add base arguments used for both Beta and GA. Args: parser: an argparse argument parser. """ resource_args.AddDatabaseResourceArg(parser, 'to execute the SQL query against') parser.add_argument( '--sql', required=True, help='The SQL query to issue to the database. Cloud Spanner SQL is ' 'described at https://cloud.google.com/spanner/docs/query-syntax') query_mode_choices = { 'NORMAL': ( 'Returns only the query result, without any information about ' 'the query plan.' ), 'PLAN': ( 'Returns only the query plan, without any result rows or ' 'execution statistics information.' ), 'PROFILE': ( 'Returns the query plan, overall execution statistics, ' 'operator-level execution statistics, along with the result rows.' ), 'WITH_STATS': ( 'Returns the overall (but not operator-level) execution statistics ' 'along with the results.' ), 'WITH_PLAN_AND_STATS': ( 'Returns the query plan, overall (but not operator-level) execution ' 'statistics, along with the results.' ), } parser.add_argument( '--query-mode', default='NORMAL', type=lambda x: x.upper(), choices=query_mode_choices, help='Mode in which the query must be processed.') parser.add_argument( '--enable-partitioned-dml', action='store_true', help='Execute DML statement using Partitioned DML') parser.add_argument( '--timeout', type=arg_parsers.Duration(), default='10m', help='Maximum time to wait for the SQL query to complete. See $ gcloud ' 'topic datetimes for information on duration formats.') msgs = apis.GetMessagesModule('spanner', 'v1') GetRequestPriorityMapper(msgs).choice_arg.AddToParser(parser) timestamp_bound_group = parser.add_argument_group( mutex=True, help='Read-only query timestamp bound. The default is --strong. See ' 'https://cloud.google.com/spanner/docs/timestamp-bounds.') timestamp_bound_group.add_argument( '--strong', action='store_true', help='Perform a strong query.') timestamp_bound_group.add_argument( '--read-timestamp', metavar='TIMESTAMP', help='Perform a query at the given timestamp.') parser.add_argument( '--database-role', help='Database role user assumes while accessing the database.') def GetRequestPriorityMapper(messages): return arg_utils.ChoiceEnumMapper( '--priority', messages.RequestOptions.PriorityValueValuesEnum, custom_mappings={ 'PRIORITY_LOW': 'low', 'PRIORITY_MEDIUM': 'medium', 'PRIORITY_HIGH': 'high', 'PRIORITY_UNSPECIFIED': 'unspecified', }, help_str='The priority for the execute SQL request.', ) @base.DefaultUniverseOnly @base.UnicodeIsSupported class Query(base.Command): """Executes a SQL query against a Cloud Spanner database.""" detailed_help = DETAILED_HELP @staticmethod def Args(parser): """See base class.""" AddBaseArgs(parser) def Run(self, args): """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. """ msgs = apis.GetMessagesModule('spanner', 'v1') request_options = msgs.RequestOptions( priority=GetRequestPriorityMapper(msgs).GetEnumForChoice(args.priority) ) read_only_options = self.ParseReadOnlyOptions(args) session = CreateSession(args, args.database_role) try: return database_sessions.ExecuteSql( args.sql, args.query_mode, session, read_only_options, request_options, args.enable_partitioned_dml, args.timeout) finally: database_sessions.Delete(session) def ParseReadOnlyOptions(self, args): """Parses the options for a read-only request from command line arguments. Args: args: Command line arguments. Returns: A ReadOnly message if the query is read-only (not DML), otherwise None. """ if QueryHasDml(args.sql): if args.IsSpecified('strong'): raise c_exceptions.InvalidArgumentException( '--strong', 'A timestamp bound cannot be specified for a DML statement.' ) if args.IsSpecified('read_timestamp'): raise c_exceptions.InvalidArgumentException( '--read-timestamp', 'A timestamp bound cannot be specified for a DML statement.' ) return None else: msgs = apis.GetMessagesModule('spanner', 'v1') if args.IsSpecified('read_timestamp'): return msgs.ReadOnly(readTimestamp=args.read_timestamp) elif args.IsSpecified('strong'): if not args.strong: raise c_exceptions.InvalidArgumentException( '--strong', '`--strong` cannot be set to false. ' 'Instead specify a different type of timestamp bound.' ) else: return msgs.ReadOnly(strong=True) else: # The default timestamp bound is strong. return msgs.ReadOnly(strong=True) def Display(self, args, result): """Displays the server response to a query. This is called higher up the stack to over-write default display behavior. What gets displayed depends on the mode in which the query was run. 'NORMAL': query result rows 'PLAN': query plan without execution statistics 'PROFILE': query result rows and the query plan with execution statistics Args: args: The arguments originally passed to the command. result: The output of the command before display. """ display_plan = ( args.query_mode == 'PLAN' or args.query_mode == 'PROFILE' or args.query_mode == 'WITH_PLAN_AND_STATS' ) display_results = args.query_mode != 'PLAN' if sql.QueryHasAggregateStats(result): sql.DisplayQueryAggregateStats(result.stats.queryStats, log.out) if display_plan: sql.DisplayQueryPlan(result, log.out) if display_results: sql.DisplayQueryResults( result, log.status if args.query_mode == 'PROFILE' else log.out )