336 lines
11 KiB
Python
336 lines
11 KiB
Python
#!/usr/bin/env python
|
|
"""Bigquery-specific NewCmd wrapper intended for CLI commands to subclass."""
|
|
|
|
import logging
|
|
import os
|
|
import pdb
|
|
import shlex
|
|
import sys
|
|
import traceback
|
|
import types
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from absl import app
|
|
from absl import flags
|
|
import googleapiclient
|
|
|
|
import bq_auth_flags
|
|
import bq_flags
|
|
import bq_utils
|
|
from gcloud_wrapper import bq_to_gcloud_command_executor
|
|
from utils import bq_error
|
|
from utils import bq_error_utils
|
|
from utils import bq_logging
|
|
from utils import bq_processor_utils
|
|
from pyglib import appcommands
|
|
|
|
FLAGS = flags.FLAGS
|
|
|
|
|
|
def _UseServiceAccount() -> bool:
|
|
return bool(
|
|
FLAGS.use_gce_service_account
|
|
or FLAGS.service_account
|
|
)
|
|
|
|
|
|
# TODO(user): This code uses more than the average amount of
|
|
# Python magic. Explain what the heck is going on throughout.
|
|
class NewCmd(appcommands.Cmd):
|
|
"""Featureful extension of appcommands.Cmd."""
|
|
|
|
def __init__(self, name: str, flag_values: flags.FlagValues) -> None:
|
|
super(NewCmd, self).__init__(name, flag_values)
|
|
run_with_args = getattr(self, 'RunWithArgs', None)
|
|
self._new_style = isinstance(run_with_args, types.MethodType)
|
|
if self._new_style:
|
|
func = run_with_args.__func__
|
|
code = func.__code__ # pylint: disable=redefined-outer-name
|
|
self._full_arg_list = list(code.co_varnames[: code.co_argcount])
|
|
# TODO(user): There might be some corner case where this
|
|
# is *not* the right way to determine bound vs. unbound method.
|
|
if isinstance(run_with_args.__self__, run_with_args.__self__.__class__):
|
|
self._full_arg_list.pop(0)
|
|
self._max_args = len(self._full_arg_list)
|
|
self._min_args = self._max_args - len(func.__defaults__ or [])
|
|
self._star_args = bool(code.co_flags & 0x04)
|
|
self._star_kwds = bool(code.co_flags & 0x08)
|
|
if self._star_args:
|
|
self._max_args = sys.maxsize
|
|
self.surface_in_shell = True
|
|
self.__doc__ = self.RunWithArgs.__doc__
|
|
elif (
|
|
hasattr(self.Run, '__func__')
|
|
and self.Run.__func__ is NewCmd.Run.__func__ # pytype: disable=attribute-error
|
|
):
|
|
raise appcommands.AppCommandsError(
|
|
'Subclasses of NewCmd must override Run or RunWithArgs'
|
|
)
|
|
|
|
def __getattr__(self, name: str):
|
|
if name in self._command_flags:
|
|
return self._command_flags[name].value
|
|
return super(NewCmd, self).__getattribute__(name)
|
|
|
|
def _GetFlag(self, flagname: str) -> Optional[flags.FlagHolder]:
|
|
if flagname in self._command_flags:
|
|
return self._command_flags[flagname]
|
|
else:
|
|
return None
|
|
|
|
def _CheckFlags(self) -> None:
|
|
"""Validate flags after command specific flags have been loaded.
|
|
|
|
This function will run through all values in appcommands._cmd_argv and
|
|
pick out any unused flags and verify their validity. If the flag is
|
|
not defined, we will print the flags.FlagsError text and exit; otherwise,
|
|
we will print a positioning error message and exit. Print statements
|
|
were used in this case because raising app.UsageError caused the usage
|
|
help text to be printed.
|
|
|
|
If no extraneous flags exist, this function will do nothing.
|
|
"""
|
|
unused_flags = [
|
|
f
|
|
for f in appcommands.GetCommandArgv()
|
|
if f.startswith('--') or f.startswith('-')
|
|
]
|
|
for flag in unused_flags:
|
|
flag_name = flag[4:] if flag.startswith('--no') else flag[2:]
|
|
flag_name = flag_name.split('=')[0]
|
|
if flag_name not in FLAGS:
|
|
print((
|
|
"FATAL Flags parsing error: Unknown command line flag '%s'\n"
|
|
"Run 'bq help' to get help" % flag
|
|
))
|
|
sys.exit(1)
|
|
else:
|
|
print((
|
|
"FATAL Flags positioning error: Flag '%s' appears after final "
|
|
'command line argument. Please reposition the flag.\n'
|
|
"Run 'bq help' to get help." % flag
|
|
))
|
|
sys.exit(1)
|
|
|
|
def Run(self, argv: List[str]) -> int:
|
|
"""Run this command.
|
|
|
|
If self is a new-style command, we set up arguments and call
|
|
self.RunWithArgs, gracefully handling exceptions. If not, we
|
|
simply call self.Run(argv).
|
|
|
|
Args:
|
|
argv: List of arguments as strings.
|
|
|
|
Returns:
|
|
0 on success, nonzero on failure.
|
|
"""
|
|
self._CheckFlags()
|
|
logging.debug('In NewCmd.Run: %s', argv)
|
|
self._debug_mode = FLAGS.debug_mode
|
|
if not self._new_style:
|
|
return super(NewCmd, self).Run(argv)
|
|
|
|
original_values = {
|
|
name: self._command_flags[name].value for name in self._command_flags
|
|
}
|
|
original_presence = {
|
|
name: self._command_flags[name].present for name in self._command_flags
|
|
}
|
|
try:
|
|
args = self._command_flags(argv)[1:]
|
|
for flag_name in self._command_flags:
|
|
value = self._command_flags[flag_name].value
|
|
setattr(self, flag_name, value)
|
|
if value == original_values[flag_name]:
|
|
original_values.pop(flag_name)
|
|
new_args = []
|
|
for argname in self._full_arg_list[: self._min_args]:
|
|
flag = self._GetFlag(argname)
|
|
if flag is not None and flag.present:
|
|
new_args.append(flag.value)
|
|
elif args:
|
|
new_args.append(args.pop(0))
|
|
else:
|
|
print('Not enough positional args, still looking for %s' % (argname,))
|
|
if self.usage:
|
|
print('Usage: %s' % (self.usage,))
|
|
return 1
|
|
|
|
new_kwds = {}
|
|
for argname in self._full_arg_list[self._min_args :]:
|
|
flag = self._GetFlag(argname)
|
|
if flag is not None and flag.present:
|
|
new_kwds[argname] = flag.value
|
|
elif args:
|
|
new_kwds[argname] = args.pop(0)
|
|
|
|
if args and not self._star_args:
|
|
print('Too many positional args, still have %s' % (args,))
|
|
return 1
|
|
new_args.extend(args)
|
|
|
|
if self._debug_mode:
|
|
return self.RunDebug(new_args, new_kwds)
|
|
else:
|
|
return self.RunSafely(new_args, new_kwds)
|
|
finally:
|
|
for flag, value in original_values.items():
|
|
setattr(self, flag, value)
|
|
self._command_flags[flag].value = value
|
|
self._command_flags[flag].present = original_presence[flag]
|
|
|
|
def RunCmdLoop(self, argv) -> int:
|
|
"""Hook for use in cmd.Cmd-based command shells."""
|
|
try:
|
|
args = shlex.split(argv)
|
|
except ValueError as e:
|
|
raise SyntaxError(bq_logging.EncodeForPrinting(e)) from e
|
|
return self.Run([self._command_name] + args)
|
|
|
|
def _HandleError(self, e):
|
|
message = bq_logging.EncodeForPrinting(e)
|
|
if isinstance(e, bq_error.BigqueryClientConfigurationError):
|
|
message += ' Try running "bq init".'
|
|
print(
|
|
'Exception raised in %s operation: %s' % (self._command_name, message)
|
|
)
|
|
return 1
|
|
|
|
def RunDebug(self, args: List[str], kwds: Dict[str, Any]) -> int:
|
|
"""Run this command in debug mode."""
|
|
logging.debug('In NewCmd.RunDebug: %s, %s', args, kwds)
|
|
try:
|
|
return_value = self.RunWithArgs(*args, **kwds)
|
|
# pylint: disable=broad-except
|
|
except (BaseException, googleapiclient.errors.ResumableUploadError) as e:
|
|
# Don't break into the debugger for expected exceptions.
|
|
if (
|
|
isinstance(e, app.UsageError)
|
|
or (
|
|
isinstance(e, bq_error.BigqueryError)
|
|
and not isinstance(e, bq_error.BigqueryInterfaceError)
|
|
)
|
|
or isinstance(e, googleapiclient.errors.ResumableUploadError)
|
|
):
|
|
return self._HandleError(e)
|
|
print()
|
|
print('****************************************************')
|
|
print('** Unexpected Exception raised in bq execution! **')
|
|
if FLAGS.headless:
|
|
print('** --headless mode enabled, exiting. **')
|
|
print('** See STDERR for traceback. **')
|
|
else:
|
|
print('** --debug_mode enabled, starting pdb. **')
|
|
print('****************************************************')
|
|
print()
|
|
traceback.print_exc()
|
|
print()
|
|
if not FLAGS.headless:
|
|
pdb.post_mortem()
|
|
return 1
|
|
return return_value
|
|
|
|
def RunSafely(self, args: List[str], kwds: Dict[str, Any]) -> int:
|
|
"""Run this command, turning exceptions into print statements."""
|
|
logging.debug('In NewCmd.RunSafely: %s, %s', args, kwds)
|
|
try:
|
|
return_value = self.RunWithArgs(*args, **kwds)
|
|
# pylint: disable=broad-exception-caught
|
|
except BaseException as e:
|
|
# pylint: enable=broad-exception-caught
|
|
return self._HandleError(e)
|
|
return return_value
|
|
|
|
|
|
class BigqueryCmd(NewCmd):
|
|
"""Bigquery-specific NewCmd wrapper."""
|
|
|
|
def _NeedsInit(self) -> bool:
|
|
"""Returns true if this command requires the init command before running.
|
|
|
|
Subclasses will override for any exceptional cases.
|
|
"""
|
|
if bq_auth_flags.USE_GOOGLE_AUTH.value:
|
|
return False
|
|
return not _UseServiceAccount() and not (
|
|
os.path.exists(bq_utils.GetBigqueryRcFilename())
|
|
or os.path.exists(FLAGS.credential_file)
|
|
)
|
|
|
|
def Run(self, argv: List[str]) -> int:
|
|
"""Bigquery commands run `init` before themselves if needed."""
|
|
|
|
if FLAGS.debug_mode:
|
|
cmd_flags = [
|
|
FLAGS[f].serialize().strip() for f in FLAGS if FLAGS[f].present
|
|
]
|
|
print(' '.join(sorted(set(f for f in cmd_flags if f))))
|
|
|
|
bq_logging.ConfigureLogging(bq_flags.APILOG.value)
|
|
logging.debug('In BigqueryCmd.Run: %s', argv)
|
|
if self._NeedsInit():
|
|
appcommands.GetCommandByName('init').Run(['init'])
|
|
return super(BigqueryCmd, self).Run(argv)
|
|
|
|
def RunSafely(self, args: List[str], kwds: Dict[str, Any]) -> int:
|
|
"""Run this command, printing information about any exceptions raised."""
|
|
logging.debug('In BigqueryCmd.RunSafely: %s, %s', args, kwds)
|
|
try:
|
|
return_value = self.RunWithArgs(*args, **kwds)
|
|
except SystemExit as e:
|
|
return_value = e.code
|
|
except BaseException as e: # pylint: disable=broad-exception-caught
|
|
return bq_error_utils.process_error(e, name=self._command_name)
|
|
return return_value
|
|
|
|
def PrintJobStartInfo(self, job) -> None:
|
|
"""Print a simple status line."""
|
|
if bq_flags.FORMAT.value in ['prettyjson', 'json']:
|
|
bq_utils.PrintFormattedJsonObject(job)
|
|
else:
|
|
reference = bq_processor_utils.ConstructObjectReference(job)
|
|
print('Successfully started %s %s' % (self._command_name, reference))
|
|
|
|
def _ProcessCommandRc(self, fv):
|
|
bq_utils.ProcessBigqueryrcSection(self._command_name, fv)
|
|
|
|
def ParseCommandFlagsSharedWithAllResources(self) -> Dict[str, str]:
|
|
"""Parses flags for the command that are shared with all resources.
|
|
|
|
This is intended to be implemented by any subclass that needs it.
|
|
|
|
Returns:
|
|
A dictionary of command flags that are shared with all resources in the
|
|
command. For example `max_results` in the list command.
|
|
"""
|
|
return {}
|
|
|
|
def PossiblyDelegateToGcloudAndExit(
|
|
self,
|
|
resource: str,
|
|
bq_command: str,
|
|
identifier: Optional[str] = None,
|
|
command_flags_for_this_resource: Optional[Dict[str, str]] = None,
|
|
):
|
|
pass # pylint: disable=unreachable
|
|
|
|
def DelegateToGcloudAndExit(
|
|
self,
|
|
resource: str,
|
|
bq_command: str,
|
|
identifier: Optional[str] = None,
|
|
command_flags_for_this_resource: Optional[Dict[str, str]] = None,
|
|
):
|
|
bq_command_flags = {
|
|
**(command_flags_for_this_resource or {}),
|
|
**self.ParseCommandFlagsSharedWithAllResources(),
|
|
}
|
|
exit_code = bq_to_gcloud_command_executor.run_bq_command_using_gcloud(
|
|
resource,
|
|
bq_command,
|
|
bq_command_flags=bq_command_flags,
|
|
identifier=identifier,
|
|
)
|
|
sys.exit(exit_code)
|