390 lines
14 KiB
Python
390 lines
14 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2017 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.
|
|
"""Common methods to display parts of SQL query results."""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import unicode_literals
|
|
|
|
from functools import partial
|
|
from apitools.base.py import encoding
|
|
from googlecloudsdk.core.resource import resource_printer
|
|
from googlecloudsdk.core.util import text
|
|
from sqlparse import lexer
|
|
from sqlparse import tokens as T
|
|
|
|
|
|
def _GetAdditionalProperty(properties, property_key, not_found_value='Unknown'):
|
|
"""Gets the value for the given key in a list of properties.
|
|
|
|
Looks through a list of properties and tries to find the value for the given
|
|
key. If it's not found, not_found_value is returned.
|
|
|
|
Args:
|
|
properties: A dictionary of key string, value string pairs.
|
|
property_key: The key string for which we want to get the value.
|
|
not_found_value: The string value to return if the key is not found.
|
|
|
|
Returns:
|
|
A string containing the value for the given key, or `not_found_value` if
|
|
the key is not found.
|
|
"""
|
|
for prop in properties:
|
|
if prop.key == property_key:
|
|
if hasattr(prop, 'value'):
|
|
return prop.value
|
|
break
|
|
return not_found_value
|
|
|
|
|
|
def _ConvertToTree(plan_nodes):
|
|
"""Creates tree of Node objects from the plan_nodes in server response.
|
|
|
|
Args:
|
|
plan_nodes (spanner_v1_messages.PlanNode[]): The plan_nodes from the server
|
|
response. Plan nodes are topologically sorted.
|
|
|
|
Returns:
|
|
A Node, root of a tree built from `plan_nodes`.
|
|
"""
|
|
# plan_nodes is a topologically sorted list, with the root node first.
|
|
return _BuildSubTree(plan_nodes, plan_nodes[0])
|
|
|
|
|
|
def _BuildSubTree(plan_nodes, node):
|
|
"""Helper for building the subtree of a query plan node.
|
|
|
|
Args:
|
|
plan_nodes (spanner_v1_messages.PlanNode[]): The plan_nodes from the server
|
|
response. Plan nodes are topologically sorted.
|
|
node (spanner_v1_messages.PlanNode): The root node of the subtree to be
|
|
built.
|
|
|
|
Returns:
|
|
A Node object.
|
|
"""
|
|
children = None
|
|
if node.childLinks:
|
|
children = [_BuildSubTree(plan_nodes, plan_nodes[link.childIndex])
|
|
for link in node.childLinks]
|
|
return Node(node, children)
|
|
|
|
|
|
def _ConvertToStringValue(prop):
|
|
"""Converts the prop to a string if it exists.
|
|
|
|
Args:
|
|
prop (object_value): The value returned from _GetAdditionalProperty.
|
|
|
|
Returns:
|
|
A string value for the given prop, or the `not_found_value` if the prop does
|
|
not exist.
|
|
"""
|
|
return getattr(prop, 'string_value', prop)
|
|
|
|
|
|
def _DisplayNumberOfRowsModified(row_count, is_exact_count, out):
|
|
"""Prints number of rows modified by a DML statement.
|
|
|
|
Args:
|
|
row_count: Either the exact number of rows modified by statement or the
|
|
lower bound of rows modified by a Partitioned DML statement.
|
|
is_exact_count: Boolean stating whether the number is the exact count.
|
|
out: Output stream to which we print.
|
|
"""
|
|
if is_exact_count:
|
|
output_str = 'Statement modified {} {}'
|
|
else:
|
|
output_str = 'Statement modified a lower bound of {} {}'
|
|
|
|
if row_count == 1:
|
|
out.Print(output_str.format(row_count, 'row'))
|
|
else:
|
|
out.Print(output_str.format(row_count, 'rows'))
|
|
|
|
|
|
def QueryHasDml(sql):
|
|
"""Determines if the sql string contains a DML query.
|
|
|
|
Args:
|
|
sql (string): The sql string entered by the user.
|
|
|
|
Returns:
|
|
A boolean.
|
|
"""
|
|
sql = sql.lstrip().lower()
|
|
tokenized = lexer.tokenize(sql)
|
|
for token in list(tokenized):
|
|
has_dml = (
|
|
token == (T.Keyword.DML, 'insert') or
|
|
token == (T.Keyword.DML, 'update') or
|
|
token == (T.Keyword.DML, 'delete'))
|
|
if has_dml:
|
|
return True
|
|
return False
|
|
|
|
|
|
def QueryHasAggregateStats(result):
|
|
"""Checks if the given results have aggregate statistics.
|
|
|
|
Args:
|
|
result (spanner_v1_messages.ResultSetStats): The stats for a query.
|
|
|
|
Returns:
|
|
A boolean indicating whether 'results' contain aggregate statistics.
|
|
"""
|
|
return hasattr(
|
|
result, 'stats') and getattr(result.stats, 'queryStats', None) is not None
|
|
|
|
|
|
def DisplayQueryAggregateStats(query_stats, out):
|
|
"""Displays the aggregate stats for a Spanner SQL query.
|
|
|
|
Looks at the queryStats portion of the query response and prints some of
|
|
the aggregate statistics.
|
|
|
|
Args:
|
|
query_stats (spanner_v1_messages.ResultSetStats.QueryStatsValue): The query
|
|
stats taken from the server response to a query.
|
|
out: Output stream to which we print.
|
|
"""
|
|
get_prop = partial(_GetAdditionalProperty, query_stats.additionalProperties)
|
|
stats = {
|
|
'total_elapsed_time': _ConvertToStringValue(get_prop('elapsed_time')),
|
|
'cpu_time': _ConvertToStringValue(get_prop('cpu_time')),
|
|
'rows_returned': _ConvertToStringValue(get_prop('rows_returned')),
|
|
'rows_scanned': _ConvertToStringValue(get_prop('rows_scanned')),
|
|
'optimizer_version': _ConvertToStringValue(get_prop('optimizer_version')),
|
|
}
|
|
resource_printer.Print(
|
|
stats,
|
|
'table[box](total_elapsed_time, cpu_time, rows_returned, rows_scanned, optimizer_version)',
|
|
out=out)
|
|
|
|
|
|
def DisplayQueryPlan(result, out):
|
|
"""Displays a graphical query plan for a query.
|
|
|
|
Args:
|
|
result (spanner_v1_messages.ResultSet): The server response to a query.
|
|
out: Output stream to which we print.
|
|
"""
|
|
node_tree_root = _ConvertToTree(result.stats.queryPlan.planNodes)
|
|
node_tree_root.PrettyPrint(out)
|
|
|
|
|
|
def DisplayQueryResults(result, out):
|
|
"""Prints the result rows for a query.
|
|
|
|
Args:
|
|
result (spanner_v1_messages.ResultSet): The server response to a query.
|
|
out: Output stream to which we print.
|
|
"""
|
|
if hasattr(result.stats,
|
|
'rowCountExact') and result.stats.rowCountExact is not None:
|
|
_DisplayNumberOfRowsModified(result.stats.rowCountExact, True, out)
|
|
|
|
if hasattr(
|
|
result.stats,
|
|
'rowCountLowerBound') and result.stats.rowCountLowerBound is not None:
|
|
_DisplayNumberOfRowsModified(result.stats.rowCountLowerBound, False, out)
|
|
|
|
if result.metadata.rowType.fields:
|
|
# Print "(Unspecified)" for computed columns.
|
|
fields = [
|
|
field.name or '(Unspecified)'
|
|
for field in result.metadata.rowType.fields
|
|
]
|
|
|
|
# Create the format string we pass to the table layout.
|
|
table_format = ','.join('row.slice({0}).join():label="{1}"'.format(i, f)
|
|
for i, f in enumerate(fields))
|
|
rows = [{
|
|
'row': encoding.MessageToPyValue(row.entry)
|
|
} for row in result.rows]
|
|
|
|
# Can't use the PrintText method because we want special formatting.
|
|
resource_printer.Print(rows, 'table({0})'.format(table_format), out=out)
|
|
|
|
|
|
class Node(object):
|
|
"""Represents a single node in a Spanner query plan.
|
|
|
|
Attributes:
|
|
properties (spanner_v1_messages.PlanNode): The details about a given node
|
|
as returned from the server.
|
|
children: A list of children in the query plan of type Node.
|
|
"""
|
|
|
|
def __init__(self, properties, children=None):
|
|
self.children = children or []
|
|
self.properties = properties
|
|
|
|
def _DisplayKindAndName(self, out, prepend, stub):
|
|
"""Prints the kind of the node (SCALAR or RELATIONAL) and its name."""
|
|
kind_and_name = '{}{} {} {}'.format(prepend, stub, self.properties.kind,
|
|
self.properties.displayName)
|
|
out.Print(kind_and_name)
|
|
|
|
def _GetNestedStatProperty(self, prop_name, nested_prop_name):
|
|
"""Gets a nested property name on this object's executionStats.
|
|
|
|
Args:
|
|
prop_name: A string of the key name for the outer property on
|
|
executionStats.
|
|
nested_prop_name: A string of the key name of the nested property.
|
|
|
|
Returns:
|
|
The string value of the nested property, or None if the outermost
|
|
property or nested property don't exist.
|
|
"""
|
|
prop = _GetAdditionalProperty(
|
|
self.properties.executionStats.additionalProperties, prop_name, '')
|
|
if not prop:
|
|
return None
|
|
|
|
nested_prop = _GetAdditionalProperty(prop.object_value.properties,
|
|
nested_prop_name, '')
|
|
if nested_prop:
|
|
return nested_prop.string_value
|
|
|
|
return None
|
|
|
|
def _DisplayExecutionStats(self, out, prepend, beneath_stub):
|
|
"""Prints the relevant execution statistics for a node.
|
|
|
|
More specifically, print out latency information and the number of
|
|
executions. This information only exists when query is run in 'PROFILE'
|
|
mode.
|
|
|
|
Args:
|
|
out: Output stream to which we print.
|
|
prepend: String that precedes any information about this node to maintain
|
|
a visible hierarchy.
|
|
beneath_stub: String that preserves the indentation of the vertical lines.
|
|
"""
|
|
if not self.properties.executionStats:
|
|
return None
|
|
|
|
stat_props = []
|
|
|
|
num_executions = self._GetNestedStatProperty('execution_summary',
|
|
'num_executions')
|
|
if num_executions:
|
|
num_executions = int(num_executions)
|
|
executions_str = '{} {}'.format(num_executions,
|
|
text.Pluralize(num_executions,
|
|
'execution'))
|
|
stat_props.append(executions_str)
|
|
|
|
# Total latency and latency unit are always expected to be present when
|
|
# latency exists. Latency exists when the query is run in PROFILE mode.
|
|
mean_latency = self._GetNestedStatProperty('latency', 'mean')
|
|
total_latency = self._GetNestedStatProperty('latency', 'total')
|
|
unit = self._GetNestedStatProperty('latency', 'unit')
|
|
if mean_latency:
|
|
stat_props.append('{} {} average latency'.format(mean_latency, unit))
|
|
elif total_latency:
|
|
stat_props.append('{} {} total latency'.format(total_latency, unit))
|
|
|
|
if stat_props:
|
|
executions_stats_str = '{}{} ({})'.format(prepend, beneath_stub,
|
|
', '.join(stat_props))
|
|
out.Print(executions_stats_str)
|
|
|
|
def _DisplayMetadata(self, out, prepend, beneath_stub):
|
|
"""Prints the keys and values of the metadata for a node.
|
|
|
|
Args:
|
|
out: Output stream to which we print.
|
|
prepend: String that precedes any information about this node to maintain
|
|
a visible hierarchy.
|
|
beneath_stub: String that preserves the indentation of the vertical lines.
|
|
"""
|
|
if self.properties.metadata:
|
|
additional_props = []
|
|
# additionalProperties looks like: [key: {value: {string_value: str}}]
|
|
for prop in self.properties.metadata.additionalProperties:
|
|
additional_props.append(
|
|
'{}: {}'.format(prop.key, prop.value.string_value))
|
|
metadata = '{}{} {}'.format(prepend, beneath_stub,
|
|
', '.join(sorted(additional_props)))
|
|
out.Print(metadata)
|
|
|
|
def _DisplayShortRepresentation(self, out, prepend, beneath_stub):
|
|
if self.properties.shortRepresentation:
|
|
short_rep = '{}{} {}'.format(
|
|
prepend, beneath_stub,
|
|
self.properties.shortRepresentation.description)
|
|
out.Print(short_rep)
|
|
|
|
def _DisplayBreakLine(self, out, prepend, beneath_stub, is_root):
|
|
"""Displays an empty line between nodes for visual breathing room.
|
|
|
|
Keeps in tact the vertical lines connecting all immediate children of a
|
|
node to each other.
|
|
|
|
Args:
|
|
out: Output stream to which we print.
|
|
prepend: String that precedes any information about this node to maintain
|
|
a visible hierarchy.
|
|
beneath_stub: String that preserves the indentation of the vertical lines.
|
|
is_root: Boolean indicating whether this node is the root of the tree.
|
|
"""
|
|
above_child = ' ' if is_root else ''
|
|
above_child += ' |' if self.children else ''
|
|
break_line = '{}{}{}'.format(prepend, beneath_stub, above_child)
|
|
# It could be the case the beneath_stub adds spaces but above_child doesn't
|
|
# add an additional vertical line, in which case we want to remove the
|
|
# extra trailing spaces.
|
|
out.Print(break_line.rstrip())
|
|
|
|
def PrettyPrint(self, out, prepend=None, is_last=True, is_root=True):
|
|
"""Prints a string representation of this node in the tree.
|
|
|
|
Args:
|
|
out: Output stream to which we print.
|
|
prepend: String that precedes any information about this node to maintain
|
|
a visible hierarchy.
|
|
is_last: Boolean indicating whether this node is the last child of its
|
|
parent.
|
|
is_root: Boolean indicating whether this node is the root of the tree.
|
|
"""
|
|
prepend = prepend or ''
|
|
# The symbol immediately before node kind to indicate that this is a child
|
|
# of its parents. All nodes except the root get one.
|
|
stub = '' if is_root else (r'\-' if is_last else '+-')
|
|
|
|
# To list additional properties beneath the name, figure out how they should
|
|
# be indented relative to the name's stub.
|
|
beneath_stub = '' if is_root else (' ' if is_last else '| ')
|
|
|
|
self._DisplayKindAndName(out, prepend, stub)
|
|
self._DisplayExecutionStats(out, prepend, beneath_stub)
|
|
self._DisplayMetadata(out, prepend, beneath_stub)
|
|
self._DisplayShortRepresentation(out, prepend, beneath_stub)
|
|
self._DisplayBreakLine(out, prepend, beneath_stub, is_root)
|
|
|
|
for idx, child in enumerate(self.children):
|
|
is_last_child = idx == len(self.children) - 1
|
|
# The amount each subsequent level in the tree is indented.
|
|
indent = ' '
|
|
# Connect all immediate children to each other with a vertical line
|
|
# of '|'. Don't extend this line down past the last child node. It's
|
|
# cleaner.
|
|
child_prepend = prepend + (' ' if is_last else '|') + indent
|
|
child.PrettyPrint(
|
|
out, prepend=child_prepend, is_last=is_last_child, is_root=False)
|