182 lines
6.3 KiB
Python
182 lines
6.3 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 to query metadata against Dataproc Metastore services database."""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import unicode_literals
|
|
|
|
import io
|
|
import json
|
|
import posixpath
|
|
|
|
from apitools.base.py import exceptions as apitools_exceptions
|
|
from googlecloudsdk.api_lib.metastore import operations_util
|
|
from googlecloudsdk.api_lib.metastore import services_util as services_api_util
|
|
from googlecloudsdk.api_lib.metastore import util as api_util
|
|
from googlecloudsdk.api_lib.storage import storage_api
|
|
from googlecloudsdk.api_lib.storage import storage_util
|
|
from googlecloudsdk.calliope import base
|
|
from googlecloudsdk.command_lib.metastore import resource_args
|
|
from googlecloudsdk.core import log
|
|
from googlecloudsdk.core.resource import resource_printer
|
|
import six
|
|
|
|
DETAILED_HELP = {
|
|
'EXAMPLES':
|
|
"""\
|
|
To query metadata against a Dataproc Metastore service with the name
|
|
`my-metastore-service` in location `us-central1`, and the sql query
|
|
"show tables;", run:
|
|
|
|
$ {command} my-metastore-service --location=us-central1
|
|
--query="show tables;"
|
|
|
|
"""
|
|
}
|
|
|
|
|
|
def AddBaseArgs(parser):
|
|
"""Parses provided arguments to add base arguments used for Alpha/Beta/GA.
|
|
|
|
Args:
|
|
parser: an argparse argument parser.
|
|
"""
|
|
resource_args.AddServiceResourceArg(
|
|
parser, 'to query metadata', plural=False, required=True, positional=True)
|
|
parser.add_argument(
|
|
'--query',
|
|
required=True,
|
|
help="""\
|
|
Use Google Standard SQL query for Cloud Spanner and MySQL query
|
|
syntax for Cloud SQL. Cloud Spanner SQL is described at
|
|
https://cloud.google.com/spanner/docs/query-syntax)"
|
|
""",
|
|
)
|
|
|
|
|
|
@base.UnicodeIsSupported
|
|
@base.ReleaseTracks(
|
|
base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA, base.ReleaseTrack.GA
|
|
)
|
|
@base.DefaultUniverseOnly
|
|
class Query(base.Command):
|
|
"""Execute a SQL query against a Dataproc Metastore Service's metadata."""
|
|
|
|
detailed_help = DETAILED_HELP
|
|
|
|
@staticmethod
|
|
def Args(parser):
|
|
"""See base class."""
|
|
AddBaseArgs(parser)
|
|
base.FORMAT_FLAG.AddToParser(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.
|
|
"""
|
|
env_ref = args.CONCEPTS.service.Parse()
|
|
operation = None
|
|
try:
|
|
operation = services_api_util.QueryMetadata(
|
|
env_ref.RelativeName(), args.query, release_track=self.ReleaseTrack())
|
|
log.out.Print('with operation [{}]'.format(operation.name))
|
|
except apitools_exceptions.HttpError:
|
|
raise api_util.QueryMetadataError('Query did not succeed.')
|
|
operation_result = None
|
|
try:
|
|
operation_result = operations_util.PollAndReturnOperation(
|
|
operation,
|
|
'Waiting for [{}] to query'.format(env_ref.RelativeName()),
|
|
release_track=self.ReleaseTrack())
|
|
except api_util.OperationError as e:
|
|
log.UpdatedResource(
|
|
env_ref.RelativeName(),
|
|
kind='service',
|
|
is_async=False,
|
|
failed=six.text_type(e))
|
|
if (operation_result is None or
|
|
not operation_result.additionalProperties or
|
|
len(operation_result.additionalProperties) < 2):
|
|
return None
|
|
result_manifest_uri = None
|
|
for message in operation_result.additionalProperties:
|
|
if message.key == 'resultManifestUri':
|
|
result_manifest_uri = message.value.string_value
|
|
if result_manifest_uri is None:
|
|
return None
|
|
gcs_client = storage_api.StorageClient()
|
|
result_manifest_json = json.load(
|
|
io.TextIOWrapper(
|
|
gcs_client.ReadObject(
|
|
storage_util.ObjectReference.FromUrl(result_manifest_uri,
|
|
True)),
|
|
encoding='utf-8'))
|
|
|
|
# Query succeed
|
|
log.out.Print(result_manifest_json['status']['message'],
|
|
result_manifest_uri)
|
|
if not result_manifest_json['filenames']:
|
|
return None
|
|
if len(result_manifest_json['filenames']) > 1:
|
|
log.out.Print('The number of rows exceeds 1000 to display. ' +
|
|
'Please find more results at the cloud storage location.')
|
|
query_result_file_name = result_manifest_json['filenames'][0]
|
|
return json.load(
|
|
io.TextIOWrapper(
|
|
gcs_client.ReadObject(
|
|
# GCS URIs follow POSIX path formatting.
|
|
storage_util.ObjectReference.FromUrl(
|
|
posixpath.join(
|
|
posixpath.dirname(result_manifest_uri),
|
|
query_result_file_name,
|
|
),
|
|
True,
|
|
)
|
|
),
|
|
encoding='utf-8',
|
|
)
|
|
)
|
|
|
|
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.
|
|
|
|
Args:
|
|
args: The arguments originally passed to the command.
|
|
result: The output of the command before display.
|
|
"""
|
|
if not result or 'metadata' not in result or 'columns' not in result[
|
|
'metadata'] or 'rows' not in result:
|
|
return
|
|
fields = [
|
|
field['name'] or '(Unspecified)'
|
|
for field in result['metadata']['columns']
|
|
]
|
|
# 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': row} 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=log.out)
|