267 lines
9.1 KiB
Python
267 lines
9.1 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.
|
|
"""Spanner migration library functions and utilities for the spanner-migration-tool binary."""
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import unicode_literals
|
|
|
|
import copy
|
|
import os
|
|
|
|
from googlecloudsdk.command_lib.util.anthos import binary_operations
|
|
from googlecloudsdk.core import exceptions as c_except
|
|
|
|
|
|
def GetEnvArgsForCommand(extra_vars=None, exclude_vars=None):
|
|
"""Return an env dict to be passed on command invocation."""
|
|
env = copy.deepcopy(os.environ)
|
|
if extra_vars:
|
|
env.update(extra_vars)
|
|
if exclude_vars:
|
|
for k in exclude_vars:
|
|
env.pop(k)
|
|
return env
|
|
|
|
|
|
class SpannerMigrationException(c_except.Error):
|
|
"""Base Exception for any errors raised by gcloud spanner migration surface."""
|
|
|
|
|
|
class SpannerMigrationWrapper(binary_operations.StreamingBinaryBackedOperation):
|
|
"""Binary operation wrapper for spanner-migration-tool commands."""
|
|
|
|
def __init__(self, **kwargs):
|
|
super(SpannerMigrationWrapper, self).__init__(
|
|
binary='spanner-migration-tool', install_if_missing=True, **kwargs)
|
|
|
|
def _ParseSchemaArgs(self,
|
|
source,
|
|
prefix=None,
|
|
source_profile=None,
|
|
target=None,
|
|
target_profile=None,
|
|
dry_run=False,
|
|
log_level=None,
|
|
project=None,
|
|
**kwargs):
|
|
""""Parse args for the schema command."""
|
|
del kwargs
|
|
exec_args = ['schema']
|
|
if source:
|
|
exec_args.extend(['--source', source])
|
|
if prefix:
|
|
exec_args.extend(['--prefix', prefix])
|
|
if source_profile:
|
|
exec_args.extend(['--source-profile', source_profile])
|
|
if target:
|
|
exec_args.extend(['--target', target])
|
|
if target_profile:
|
|
exec_args.extend(['--target-profile', target_profile])
|
|
if dry_run:
|
|
exec_args.append('--dry-run')
|
|
if log_level:
|
|
exec_args.extend(['--log-level', log_level])
|
|
if project:
|
|
exec_args.extend(['--project', project])
|
|
return exec_args
|
|
|
|
def _ParseDataArgs(self,
|
|
source,
|
|
session,
|
|
prefix=None,
|
|
skip_foreign_keys=False,
|
|
source_profile=None,
|
|
target=None,
|
|
target_profile=None,
|
|
write_limit=None,
|
|
dry_run=False,
|
|
log_level=None,
|
|
project=None,
|
|
dataflow_template=None,
|
|
**kwargs):
|
|
""""Parse args for the data command."""
|
|
del kwargs
|
|
exec_args = ['data']
|
|
if source:
|
|
exec_args.extend(['--source', source])
|
|
if session:
|
|
exec_args.extend(['--session', session])
|
|
if prefix:
|
|
exec_args.extend(['--prefix', prefix])
|
|
if skip_foreign_keys:
|
|
exec_args.append('--skip-foreign-keys')
|
|
if source_profile:
|
|
exec_args.extend(['--source-profile', source_profile])
|
|
if target:
|
|
exec_args.extend(['--target', target])
|
|
if target_profile:
|
|
exec_args.extend(['--target-profile', target_profile])
|
|
if write_limit:
|
|
exec_args.extend(['--write-limit', write_limit])
|
|
if dry_run:
|
|
exec_args.append('--dry-run')
|
|
if log_level:
|
|
exec_args.extend(['--log-level', log_level])
|
|
if project:
|
|
exec_args.extend(['--project', project])
|
|
if dataflow_template:
|
|
exec_args.extend(['--dataflow-template', dataflow_template])
|
|
return exec_args
|
|
|
|
def _ParseSchemaAndDataArgs(self,
|
|
source,
|
|
prefix=None,
|
|
skip_foreign_keys=False,
|
|
source_profile=None,
|
|
target=None,
|
|
target_profile=None,
|
|
write_limit=None,
|
|
dry_run=False,
|
|
log_level=None,
|
|
project=None,
|
|
dataflow_template=None,
|
|
**kwargs):
|
|
""""Parse args for the schema-and-data command."""
|
|
del kwargs
|
|
exec_args = ['schema-and-data']
|
|
if source:
|
|
exec_args.extend(['--source', source])
|
|
if prefix:
|
|
exec_args.extend(['--prefix', prefix])
|
|
if skip_foreign_keys:
|
|
exec_args.append('--skip-foreign-keys')
|
|
if source_profile:
|
|
exec_args.extend(['--source-profile', source_profile])
|
|
if target:
|
|
exec_args.extend(['--target', target])
|
|
if target_profile:
|
|
exec_args.extend(['--target-profile', target_profile])
|
|
if write_limit:
|
|
exec_args.extend(['--write-limit', write_limit])
|
|
if dry_run:
|
|
exec_args.append('--dry-run')
|
|
if log_level:
|
|
exec_args.extend(['--log-level', log_level])
|
|
if project:
|
|
exec_args.extend(['--project', project])
|
|
if dataflow_template:
|
|
exec_args.extend(['--dataflow-template', dataflow_template])
|
|
return exec_args
|
|
|
|
def _ParseWebArgs(self,
|
|
open_flag=False,
|
|
port=None,
|
|
log_level=None,
|
|
dataflow_template=None,
|
|
**kwargs):
|
|
"""Parse args for the web command."""
|
|
del kwargs
|
|
exec_args = ['web']
|
|
if open_flag:
|
|
exec_args.append('--open')
|
|
if port:
|
|
exec_args.extend(['--port', port])
|
|
if log_level:
|
|
exec_args.extend(['--log-level', log_level])
|
|
if dataflow_template:
|
|
exec_args.extend(['--dataflow-template', dataflow_template])
|
|
return exec_args
|
|
|
|
def ParseCleanupArgs(self,
|
|
job_id,
|
|
data_shard_ids=None,
|
|
target_profile=None,
|
|
datastream=False,
|
|
dataflow=False,
|
|
pub_sub=False,
|
|
monitoring=False,
|
|
log_level=None,
|
|
**kwargs):
|
|
""""Parse args for the cleanup command."""
|
|
del kwargs
|
|
exec_args = ['cleanup']
|
|
if job_id:
|
|
exec_args.extend(['--jobId', job_id])
|
|
if data_shard_ids:
|
|
exec_args.extend(['--dataShardIds', data_shard_ids])
|
|
if target_profile:
|
|
exec_args.extend(['--target-profile', target_profile])
|
|
if datastream:
|
|
exec_args.append('--datastream')
|
|
if dataflow:
|
|
exec_args.append('--dataflow')
|
|
if pub_sub:
|
|
exec_args.append('--pubsub')
|
|
if monitoring:
|
|
exec_args.append('--monitoring')
|
|
if log_level:
|
|
exec_args.append('--log-level')
|
|
return exec_args
|
|
|
|
def ParseImportArgs(self,
|
|
instance,
|
|
database,
|
|
source_uri,
|
|
source_format,
|
|
table_name=None,
|
|
project=None,
|
|
schema_uri=None,
|
|
csv_line_delimiter=None,
|
|
csv_field_delimiter=None,
|
|
database_dialect=None,
|
|
**kwargs):
|
|
""""Parse args for the import command."""
|
|
del kwargs
|
|
exec_args = ['import']
|
|
if instance:
|
|
exec_args.extend(['--instance', instance])
|
|
if database:
|
|
exec_args.extend(['--database', database])
|
|
if table_name:
|
|
exec_args.extend(['--table-name', table_name])
|
|
if source_uri:
|
|
exec_args.extend(['--source-uri', source_uri])
|
|
if source_format:
|
|
exec_args.extend(['--source-format', source_format])
|
|
if schema_uri:
|
|
exec_args.extend(['--schema-uri', schema_uri])
|
|
if csv_line_delimiter:
|
|
exec_args.extend(['--csv-line-delimiter', csv_line_delimiter])
|
|
if csv_field_delimiter:
|
|
exec_args.extend(['--csv-field-delimiter', csv_field_delimiter])
|
|
if project:
|
|
exec_args.extend(['--project', project])
|
|
if database_dialect:
|
|
exec_args.extend(['--database-dialect', database_dialect])
|
|
return exec_args
|
|
|
|
def _ParseArgsForCommand(self, command, **kwargs):
|
|
"""Call the parser corresponding to the command."""
|
|
if command == 'schema':
|
|
return self._ParseSchemaArgs(**kwargs)
|
|
elif command == 'data':
|
|
return self._ParseDataArgs(**kwargs)
|
|
elif command == 'schema-and-data':
|
|
return self._ParseSchemaAndDataArgs(**kwargs)
|
|
elif command == 'web':
|
|
return self._ParseWebArgs(**kwargs)
|
|
elif command == 'cleanup':
|
|
return self.ParseCleanupArgs(**kwargs)
|
|
elif command == 'import':
|
|
return self.ParseImportArgs(**kwargs)
|
|
else:
|
|
raise binary_operations.InvalidOperationForBinary(
|
|
'Invalid Operation [{}] for spanner-migration-tool'.format(command))
|