276 lines
8.4 KiB
Python
276 lines
8.4 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2018 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.
|
|
"""General BQ surface command utilites for python commands."""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
from apitools.base.py import encoding
|
|
from googlecloudsdk.api_lib.util import waiter
|
|
from googlecloudsdk.calliope import exceptions
|
|
from googlecloudsdk.command_lib.util.apis import arg_utils
|
|
from googlecloudsdk.command_lib.util.apis import yaml_data
|
|
from googlecloudsdk.command_lib.util.args import resource_args
|
|
from googlecloudsdk.command_lib.util.concepts import concept_parsers
|
|
from googlecloudsdk.core import yaml
|
|
import six
|
|
|
|
|
|
DEFAULT_MAX_QUERY_RESULTS = 1000
|
|
|
|
|
|
class BqJobPoller(waiter.OperationPoller):
|
|
"""Poller for managing Bq Jobs."""
|
|
|
|
def __init__(
|
|
self,
|
|
job_service,
|
|
result_service,
|
|
max_query_results=DEFAULT_MAX_QUERY_RESULTS,
|
|
):
|
|
"""Sets up poller for BigQuery Jobs.
|
|
|
|
Args:
|
|
job_service: apitools.base.py.base_api.BaseApiService, api service for
|
|
retrieving information about ongoing job.
|
|
result_service: apitools.base.py.base_api.BaseApiService, api service for
|
|
retrieving created result of initiated operation.
|
|
max_query_results: maximum number of records to return from a query job.
|
|
"""
|
|
self.result_service = result_service
|
|
self.job_service = job_service
|
|
self.max_query_results = max_query_results
|
|
|
|
def IsDone(self, job):
|
|
"""Overrides."""
|
|
if job.status.state == 'DONE':
|
|
if job.status.errorResult:
|
|
raise waiter.OperationError(job.status.errorResult.message)
|
|
return True
|
|
return False
|
|
|
|
def Poll(self, job_ref):
|
|
"""Overrides.
|
|
|
|
Args:
|
|
job_ref: googlecloudsdk.core.resources.Resource.
|
|
|
|
Returns:
|
|
fetched job message.
|
|
"""
|
|
request_type = self.job_service.GetRequestType('Get')
|
|
return self.job_service.Get(
|
|
request_type(jobId=job_ref.Name(), projectId=job_ref.Parent().Name())
|
|
)
|
|
|
|
def GetResult(self, job):
|
|
"""Overrides to get the response from the completed job by job type.
|
|
|
|
Args:
|
|
job: api_name_messages.Job.
|
|
|
|
Returns:
|
|
the 'response' field of the job.
|
|
"""
|
|
request_type = self.result_service.GetRequestType('Get')
|
|
job_type = job.configuration.jobType
|
|
|
|
if job_type == 'COPY':
|
|
result_table = job.configuration.copy.destinationTable
|
|
request = request_type(
|
|
datasetId=result_table.datasetId,
|
|
tableId=result_table.tableId,
|
|
projectId=result_table.projectId,
|
|
)
|
|
elif job_type == 'LOAD':
|
|
result_table = job.configuration.load.destinationTable
|
|
request = request_type(
|
|
datasetId=result_table.datasetId,
|
|
tableId=result_table.tableId,
|
|
projectId=result_table.projectId,
|
|
)
|
|
elif job_type == 'QUERY':
|
|
request_type = self.result_service.GetRequestType('GetQueryResults')
|
|
request = request_type(
|
|
jobId=job.jobReference.jobId,
|
|
maxResults=self.max_query_results,
|
|
projectId=job.jobReference.projectId,
|
|
)
|
|
return self.result_service.GetQueryResults(request)
|
|
else: # EXTRACT OR UNKNOWN type
|
|
return job
|
|
|
|
return self.result_service.Get(request)
|
|
|
|
|
|
class BqMigrationWorkflowPoller(waiter.OperationPoller):
|
|
"""Poller for managing BigQuery Migration Workflows."""
|
|
|
|
def __init__(
|
|
self,
|
|
migration_service,
|
|
):
|
|
"""Sets up poller for generic long running processes.
|
|
|
|
Args:
|
|
migration_service: apitools.base.py.base_api.BaseApiService, api service
|
|
for retrieving information about migration workflows.
|
|
"""
|
|
self.migration_service = migration_service
|
|
|
|
def IsDone(self, migration_workflow):
|
|
"""Overrides."""
|
|
return str(migration_workflow.state) == 'COMPLETED'
|
|
|
|
def Poll(self, migration_workflow_ref):
|
|
"""Overrides.
|
|
|
|
Args:
|
|
migration_workflow_ref: googlecloudsdk.core.resources.Resource.
|
|
|
|
Returns:
|
|
fetched migration workflow message.
|
|
"""
|
|
request_type = self.migration_service.GetRequestType('Get')
|
|
request = request_type(name=migration_workflow_ref.RelativeName())
|
|
res = self.migration_service.Get(request)
|
|
return res
|
|
|
|
def GetResult(self, migration_workflow):
|
|
"""Overrides to get the response from the completed job by job type.
|
|
|
|
Args:
|
|
migration_workflow: api_name_messages.MigrationWorkflow.
|
|
|
|
Returns:
|
|
the 'response' field of the Operation.
|
|
"""
|
|
request_type = self.migration_service.GetRequestType('Get')
|
|
request = request_type(name=migration_workflow.name)
|
|
return self.migration_service.Get(request)
|
|
|
|
|
|
def GetResourceFromFile(file_path, resource_message_type):
|
|
"""Returns the resource message and update fields in file."""
|
|
try:
|
|
resource_to_parse = yaml.load_path(file_path)
|
|
except yaml.YAMLParseError as e:
|
|
raise exceptions.BadFileException(
|
|
'File [{0}] cannot be parsed. {1}'.format(file_path, six.text_type(e))
|
|
)
|
|
except yaml.FileLoadError as e:
|
|
raise exceptions.BadFileException(
|
|
'File [{0}] cannot be opened or read. {1}'.format(
|
|
file_path, six.text_type(e)
|
|
)
|
|
)
|
|
|
|
if not isinstance(resource_to_parse, dict):
|
|
raise exceptions.BadFileException(
|
|
'File [{0}] is not a properly formatted YAML or JSON file.'.format(
|
|
file_path
|
|
)
|
|
)
|
|
|
|
try:
|
|
resource = encoding.PyValueToMessage(
|
|
resource_message_type, resource_to_parse
|
|
)
|
|
except AttributeError as e:
|
|
raise exceptions.BadFileException(
|
|
'File [{0}] is not a properly formatted YAML or JSON file. {1}'.format(
|
|
file_path, six.text_type(e)
|
|
)
|
|
)
|
|
|
|
return resource
|
|
|
|
|
|
def ProcessTableCopyOverwrite(ref, args, request):
|
|
"""Process the overwrite flag on tables copy."""
|
|
del ref # Unused
|
|
if args.overwrite:
|
|
request.job.configuration.copy.writeDisposition = 'WRITE_TRUNCATE'
|
|
return request
|
|
|
|
|
|
def ProcessTableCopyConfiguration(ref, args, request):
|
|
"""Build JobConfigurationTableCopy from request resource args."""
|
|
del ref # Unused
|
|
source_ref = args.CONCEPTS.source.Parse()
|
|
destination_ref = args.CONCEPTS.destination.Parse()
|
|
arg_utils.SetFieldInMessage(
|
|
request,
|
|
'job.configuration.copy.destinationTable.datasetId',
|
|
destination_ref.Parent().Name(),
|
|
)
|
|
arg_utils.SetFieldInMessage(
|
|
request,
|
|
'job.configuration.copy.destinationTable.projectId',
|
|
destination_ref.projectId,
|
|
)
|
|
arg_utils.SetFieldInMessage(
|
|
request,
|
|
'job.configuration.copy.destinationTable.tableId',
|
|
destination_ref.Name(),
|
|
)
|
|
arg_utils.SetFieldInMessage(
|
|
request,
|
|
'job.configuration.copy.sourceTable.datasetId',
|
|
source_ref.Parent().Name(),
|
|
)
|
|
arg_utils.SetFieldInMessage(
|
|
request,
|
|
'job.configuration.copy.sourceTable.projectId',
|
|
source_ref.projectId,
|
|
)
|
|
arg_utils.SetFieldInMessage(
|
|
request, 'job.configuration.copy.sourceTable.tableId', source_ref.Name()
|
|
)
|
|
return request
|
|
|
|
|
|
# Resource Argument utils
|
|
def GetTableCopyResourceArgs():
|
|
"""Get Table resource args (source, destination) for copy command."""
|
|
table_spec_data = yaml_data.ResourceYAMLData.FromPath('bq.table')
|
|
arg_specs = [
|
|
resource_args.GetResourcePresentationSpec(
|
|
verb='to copy from',
|
|
name='source',
|
|
required=True,
|
|
prefixes=True,
|
|
attribute_overrides={'table': 'source'},
|
|
positional=False,
|
|
resource_data=table_spec_data.GetData(),
|
|
),
|
|
resource_args.GetResourcePresentationSpec(
|
|
verb='to copy to',
|
|
name='destination',
|
|
required=True,
|
|
prefixes=True,
|
|
attribute_overrides={'table': 'destination'},
|
|
positional=False,
|
|
resource_data=table_spec_data.GetData(),
|
|
),
|
|
]
|
|
fallthroughs = {
|
|
'--source.dataset': ['--destination.dataset'],
|
|
'--destination.dataset': ['--source.dataset'],
|
|
}
|
|
return [concept_parsers.ConceptParser(arg_specs, fallthroughs)]
|