def test_parse_args_1(self):
ap = argparse.ArgumentParser()
export.add_argparse_args(ap)
args = ap.parse_args([
'--gcs-destination',
'gs://b/o',
'--content-types',
'RESOURCE, IAM_POLICY',
'--parent',
'projects/projectid',
])
self.assertIsNone(args.asset_types)
self.assertEqual(args.gcs_destination, 'gs://b/o')
self.assertEqual(args.content_types, ['RESOURCE', 'IAM_POLICY'])
self.assertEqual(args.parent, 'projects/projectid')
def test_parse_args_2(self):
ap = argparse.ArgumentParser()
export.add_argparse_args(ap)
args = ap.parse_args([
'--gcs-destination',
'gs://b/o',
'--asset-types',
'google.compute.Instance, google.compute.Firewall',
'--parent',
'projects/projectid',
])
self.assertEqual(args.gcs_destination, 'gs://b/o')
self.assertEqual(args.content_types, ['RESOURCE', 'IAM_POLICY'])
self.assertEqual(
args.asset_types,
['google.compute.Instance', 'google.compute.Firewall'])
self.assertEqual(args.parent, 'projects/projectid')
def parse_args():
"""Parse command line argments.
Present a slightly simpler interface by constructing the pipeline input and
stage parameters from the export gcs-destination if not supplied. Also
accepts arbitraty beam pipeline arguments if using a beam runner. but #
getting them into the help text is near impossible.
Returns:
List if strings. The user supplied command line arguments with
added input and stage arguments.
"""
parser = argparse.ArgumentParser()
# Take all the arguments that export takes.
export.add_argparse_args(parser, True)
parser.add_argument(
'--group_by',
default='ASSET_TYPE',
choices=['ASSET_TYPE', 'ASSET_TYPE_VERSION'],
help='How to group exported resources into Bigquery tables.')
parser.add_argument('--write_disposition',
default='WRITE_APPEND',
choices=['WRITE_APPEND', 'WRITE_EMPTY'],
help='Location to write data and load from from.')
parser.add_argument(
'--stage',
help=('Location to write intermediary data to load from from. '
'Will be --gcs-destination + "/stage" if not supplied.'))
parser.add_argument('--load_time',
default=datetime.datetime.now().isoformat(),
help=('Load time of the data (YYYY-MM-DD[HH:MM:SS])). '
'Defaults to "now".'))
parser.add_argument(
'--dataset',
help='BigQuery dataset to load to.',
required=True,
)
parser.add_argument(
'--skip-export',
help=('Do not perform asset export to GCS. Imports to bigquery'
' what\'s in the --gcs-destination. '),
action='store_true',
default=False)
parser.add_argument(
'--template-job-launch-location',
help=(
'The dataflow template to launch to import assets.'
' Should be a GCS location like '
# pylint: disable=line-too-long
'gs://professional-services-tools-asset-inventory/latest/import_pipeline'
))
parser.add_argument(
'--template-job-project',
help=('When launching a template via --template-job-launch-location, '
'this is the project id to run the job in.'))
parser.add_argument(
'--template-job-region',
default='us-central1',
help=('When launching a template via --template-job-launch-location, '
'This is the region to run in. (defaults to us-central1)'))
def json_value(string_value):
return json.loads(string_value)
parser.add_argument(
'--template-job-runtime-environment-json',
type=json_value,
help=
('When launching a template via --template-job-launch-location, '
'this is an optional json dict for '
'runtime environment for the dataflow template launch request. '
'See https://cloud.google.com/dataflow/docs/reference/rest/v1b3/RuntimeEnvironment. '
'For example : \'{"maxWorkers": 10}\''))
args, beam_args = parser.parse_known_args()
# If input isn't supplied, we can infer it from the export destination.
if 'input' not in args or not args.input:
args.input = '{}/*.json'.format(args.gcs_destination)
# If stage isn't supplied, we can infer it from export desitnation.
if 'stage' not in args or not args.stage:
args.stage = args.gcs_destination + '/stage'
return args, beam_args