def test_get_add_datasets(self):
source_dataset1_1 = BigQueryDataset(dataset_id='test1',
friendly_name='test_friendly_name',
description='test_description',
default_table_expiration_ms=24 *
60 * 60 * 1000,
location='US')
source_dataset1_2 = BigQueryDataset(dataset_id='test2',
friendly_name='foo_bar',
description='fizz_buzz',
default_table_expiration_ms=60 *
60 * 1000,
location='EU')
target_dataset1_1 = BigQueryDataset(dataset_id='test1',
friendly_name='test_friendly_name',
description='test_description',
default_table_expiration_ms=24 *
60 * 60 * 1000,
location='US')
target_dataset1_2 = BigQueryDataset(dataset_id='test2',
friendly_name='foo_bar',
description='fizz_buzz',
default_table_expiration_ms=60 *
60 * 1000,
location='EU')
source1 = [source_dataset1_1]
target1 = [target_dataset1_1, target_dataset1_2]
actual_count1, actual_results1 = DatasetAction.get_add_datasets(
source1, target1)
self.assertEqual(actual_count1, 1)
self.assertEqual(actual_results1, (target_dataset1_2, ))
source2 = [source_dataset1_1, source_dataset1_2]
target2 = [target_dataset1_1, target_dataset1_2]
actual_count2, actual_results2 = DatasetAction.get_add_datasets(
source2, target2)
self.assertEqual(actual_count2, 0)
self.assertEqual(actual_results2, ())
source3 = [source_dataset1_1, source_dataset1_2]
target3 = [target_dataset1_1]
actual_count3, actual_results3 = DatasetAction.get_add_datasets(
source3, target3)
self.assertEqual(actual_count3, 0)
self.assertEqual(actual_results3, ())
def export(ctx, output_dir, dataset, exclude_dataset):
with ThreadPoolExecutor(max_workers=ctx.obj['parallelism']) as e:
action = DatasetAction(e, project=ctx.obj['project'],
credential_file=ctx.obj['credential_file'],
no_color=not ctx.obj['color'],
debug=ctx.obj['debug'])
datasets = as_completed(action.export(output_dir, dataset, exclude_dataset))
fs = []
for d in datasets:
action = TableAction(e, d.dataset_id,
project=ctx.obj['project'],
credential_file=ctx.obj['credential_file'],
no_color=not ctx.obj['color'],
debug=ctx.obj['debug'])
fs.extend(action.export(output_dir))
as_completed(fs)
def apply_destroy(ctx, conf_dir, auto_approve, dataset, exclude_dataset):
# TODO Impl auto-approve option
destroy_counts = []
with ThreadPoolExecutor(max_workers=ctx.obj['parallelism']) as e:
dataset_action = DatasetAction(e, project=ctx.obj['project'],
credential_file=ctx.obj['credential_file'],
no_color=not ctx.obj['color'],
debug=ctx.obj['debug'])
source_datasets = [d for d in as_completed(dataset_action.list_datasets(
dataset, exclude_dataset)) if d]
target_datasets = list_local_datasets(conf_dir, dataset, exclude_dataset)
echo('------------------------------------------------------------------------')
echo()
fs = []
for d in target_datasets:
table_action = TableAction(e, d.dataset_id,
project=ctx.obj['project'],
credential_file=ctx.obj['credential_file'],
no_color=not ctx.obj['color'],
debug=ctx.obj['debug'])
source_tables = [t for t in as_completed(table_action.list_tables()) if t]
if source_tables:
echo('------------------------------------------------------------------------')
echo()
destroy_count, destroy_fs = table_action.destroy(source_tables, [])
destroy_counts.append(destroy_count)
fs.extend(destroy_fs)
as_completed(fs)
fs = []
destroy_count, destroy_fs = dataset_action.intersection_destroy(
source_datasets, target_datasets)
destroy_counts.append(destroy_count)
fs.extend(destroy_fs)
as_completed(fs)
if not any(destroy_counts):
echo(msg.MESSAGE_SUMMARY_NO_CHANGE)
echo()
else:
echo(msg.MESSAGE_APPLY_DESTROY_SUMMARY.format(sum(destroy_counts)))
echo()
def apply(ctx, conf_dir, auto_approve, dataset, exclude_dataset, mode, backup_dataset):
# TODO Impl auto-approve option
add_counts, change_counts, destroy_counts = [], [], []
with ThreadPoolExecutor(max_workers=ctx.obj['parallelism']) as e:
dataset_action = DatasetAction(e, project=ctx.obj['project'],
credential_file=ctx.obj['credential_file'],
no_color=not ctx.obj['color'],
debug=ctx.obj['debug'])
source_datasets = [d for d in as_completed(dataset_action.list_datasets(
dataset, exclude_dataset)) if d]
target_datasets = list_local_datasets(conf_dir, dataset, exclude_dataset)
echo('------------------------------------------------------------------------')
echo()
fs = []
add_count, add_fs = dataset_action.add(source_datasets, target_datasets)
add_counts.append(add_count)
fs.extend(add_fs)
change_count, change_fs = dataset_action.change(source_datasets, target_datasets)
change_counts.append(change_count)
fs.extend(change_fs)
destroy_count, destroy_fs = dataset_action.destroy(source_datasets, target_datasets)
destroy_counts.append(destroy_count)
fs.extend(destroy_fs)
as_completed(fs)
fs = []
for d in target_datasets:
target_tables = list_local_tables(conf_dir, d.dataset_id)
if target_tables is None:
continue
table_action = TableAction(e, d.dataset_id,
migration_mode=mode,
backup_dataset_id=backup_dataset,
project=ctx.obj['project'],
credential_file=ctx.obj['credential_file'],
no_color=not ctx.obj['color'],
debug=ctx.obj['debug'])
source_tables = [t for t in as_completed(table_action.list_tables()) if t]
if target_tables or source_tables:
echo('------------------------------------------------------------------------')
echo()
add_count, add_fs = table_action.add(source_tables, target_tables)
add_counts.append(add_count)
fs.extend(add_fs)
change_count, change_fs = table_action.change(source_tables, target_tables)
change_counts.append(change_count)
fs.extend(change_fs)
destroy_count, destroy_fs = table_action.destroy(source_tables, target_tables)
destroy_counts.append(destroy_count)
fs.extend(destroy_fs)
as_completed(fs)
if not any(chain.from_iterable([add_counts, change_counts, destroy_counts])):
echo(msg.MESSAGE_SUMMARY_NO_CHANGE)
echo()
else:
echo(msg.MESSAGE_APPLY_SUMMARY.format(
sum(add_counts), sum(change_counts), sum(destroy_counts)))
echo()
def plan_destroy(ctx, conf_dir, detailed_exitcode, dataset, exclude_dataset):
echo(msg.MESSAGE_PLAN_HEADER)
destroy_counts = []
with ThreadPoolExecutor(max_workers=ctx.obj['parallelism']) as e:
dataset_action = DatasetAction(e, project=ctx.obj['project'],
credential_file=ctx.obj['credential_file'],
no_color=not ctx.obj['color'],
debug=ctx.obj['debug'])
source_datasets = [d for d in as_completed(dataset_action.list_datasets(
dataset, exclude_dataset)) if d]
target_datasets = list_local_datasets(conf_dir, dataset, exclude_dataset)
echo('------------------------------------------------------------------------')
echo()
destroy_counts.append(dataset_action.plan_intersection_destroy(
source_datasets, target_datasets))
for d in target_datasets:
table_action = TableAction(e, d.dataset_id,
project=ctx.obj['project'],
credential_file=ctx.obj['credential_file'],
no_color=not ctx.obj['color'],
debug=ctx.obj['debug'])
source_tables = [t for t in as_completed(table_action.list_tables()) if t]
if source_tables:
echo('------------------------------------------------------------------------')
echo()
destroy_counts.append(table_action.plan_destroy(source_tables, []))
if not any(destroy_counts):
echo(msg.MESSAGE_SUMMARY_NO_CHANGE)
echo()
else:
echo(msg.MESSAGE_PLAN_DESTROY_SUMMARY.format(sum(destroy_counts)))
echo()
if detailed_exitcode:
sys.exit(2)
def test_get_change_datasets(self):
source_dataset1_1 = BigQueryDataset(dataset_id='test1',
friendly_name='test_friendly_name',
description='test_description',
default_table_expiration_ms=24 *
60 * 60 * 1000,
location='US')
source_dataset1_2 = BigQueryDataset(
dataset_id='test2',
friendly_name='test_friendly_name',
description='test_description',
default_table_expiration_ms=24 * 60 * 60 * 1000,
location='US',
access_entries=(BigQueryAccessEntry('OWNER', 'specialGroup',
'projectOwners'), ))
source_dataset1_3 = BigQueryDataset(dataset_id='test3',
friendly_name='test_friendly_name',
description='test_description',
default_table_expiration_ms=24 *
30 * 60 * 1000,
location='US',
labels={'foo': 'bar'})
target_dataset1_1 = BigQueryDataset(dataset_id='test1',
friendly_name='foo',
description='bar',
default_table_expiration_ms=60 *
60 * 1000,
location='UK')
target_dataset1_2 = BigQueryDataset(
dataset_id='test2',
friendly_name='test_friendly_name',
description='test_description',
default_table_expiration_ms=24 * 60 * 60 * 1000,
location='US',
access_entries=(BigQueryAccessEntry('OWNER', 'specialGroup',
'projectOwners'),
BigQueryAccessEntry(
None, 'view', {
'datasetId': 'test',
'projectId': 'test',
'tableId': 'test_table'
})))
target_dataset1_3 = BigQueryDataset(dataset_id='test3',
friendly_name='test_friendly_name',
description='test_description',
default_table_expiration_ms=24 *
60 * 60 * 1000,
location='US',
labels={
'aaa': 'bbb',
'ccc': 'ddd'
})
source1 = [source_dataset1_1, source_dataset1_2, source_dataset1_3]
target1 = [target_dataset1_1, target_dataset1_2, target_dataset1_3]
actual_count1, actual_results1 = DatasetAction.get_change_datasets(
source1, target1)
self.assertEqual(actual_count1, 3)
self.assertEqual(
set(actual_results1),
{target_dataset1_1, target_dataset1_2, target_dataset1_3})
source2 = [source_dataset1_3, source_dataset1_2, source_dataset1_1]
target2 = [target_dataset1_1, target_dataset1_2, target_dataset1_3]
actual_count2, actual_results2 = DatasetAction.get_change_datasets(
source2, target2)
self.assertEqual(actual_count2, 3)
self.assertEqual(
set(actual_results2),
{target_dataset1_1, target_dataset1_2, target_dataset1_3})
source3 = [source_dataset1_1, source_dataset1_2, source_dataset1_3]
target3 = [source_dataset1_1, source_dataset1_2, source_dataset1_3]
actual_count3, actual_results3 = DatasetAction.get_change_datasets(
source3, target3)
self.assertEqual(actual_count3, 0)
self.assertEqual(actual_results3, ())
source4 = [source_dataset1_1, source_dataset1_2, source_dataset1_3]
target4 = [source_dataset1_3, source_dataset1_2, source_dataset1_1]
actual_count4, actual_results4 = DatasetAction.get_change_datasets(
source4, target4)
self.assertEqual(actual_count4, 0)
self.assertEqual(actual_results4, ())
source5 = [target_dataset1_1, target_dataset1_2, target_dataset1_3]
target5 = [target_dataset1_1, target_dataset1_2, target_dataset1_3]
actual_count5, actual_results5 = DatasetAction.get_change_datasets(
source5, target5)
self.assertEqual(actual_count5, 0)
self.assertEqual(actual_results5, ())
source6 = [target_dataset1_1, target_dataset1_2, target_dataset1_3]
target6 = [target_dataset1_3, target_dataset1_2, target_dataset1_1]
actual_count6, actual_results6 = DatasetAction.get_change_datasets(
source6, target6)
self.assertEqual(actual_count6, 0)
self.assertEqual(actual_results6, ())
source7 = [source_dataset1_1, source_dataset1_2, source_dataset1_3]
target7 = [target_dataset1_1, target_dataset1_2]
actual_count7, actual_results7 = DatasetAction.get_change_datasets(
source7, target7)
self.assertEqual(actual_count7, 2)
self.assertEqual(set(actual_results7),
{target_dataset1_1, target_dataset1_2})
source8 = [source_dataset1_1, source_dataset1_2]
target8 = [target_dataset1_1, target_dataset1_2, target_dataset1_3]
actual_count8, actual_results8 = DatasetAction.get_change_datasets(
source8, target8)
self.assertEqual(actual_count8, 2)
self.assertEqual(set(actual_results8),
{target_dataset1_1, target_dataset1_2})
def test_get_intersection_datasets(self):
source_dataset1_1 = BigQueryDataset(dataset_id='test1',
friendly_name='test_friendly_name',
description='test_description',
default_table_expiration_ms=24 *
60 * 60 * 1000,
location='US')
source_dataset1_2 = BigQueryDataset(dataset_id='test2',
friendly_name='foo_bar',
description='fizz_buzz',
default_table_expiration_ms=60 *
60 * 1000,
location='EU')
target_dataset1_1 = BigQueryDataset(dataset_id='test1',
friendly_name='test_friendly_name',
description='test_description',
default_table_expiration_ms=24 *
60 * 60 * 1000,
location='US')
target_dataset1_2 = BigQueryDataset(dataset_id='test2',
friendly_name='foo_bar',
description='fizz_buzz',
default_table_expiration_ms=60 *
60 * 1000,
location='EU')
source1 = [source_dataset1_1]
target1 = [target_dataset1_1, target_dataset1_2]
actual_count1, actual_results1 = DatasetAction.get_intersection_datasets(
source1, target1)
self.assertEqual(actual_count1, 1)
self.assertEqual(actual_results1, (source_dataset1_1, ))
source2 = [source_dataset1_1, source_dataset1_2]
target2 = [target_dataset1_1, target_dataset1_2]
actual_count2, actual_results2 = DatasetAction.get_intersection_datasets(
source2, target2)
self.assertEqual(actual_count2, 2)
self.assertEqual(set(actual_results2),
{source_dataset1_1, source_dataset1_2})
source3 = [source_dataset1_1, source_dataset1_2]
target3 = [target_dataset1_1]
actual_count3, actual_results3 = DatasetAction.get_intersection_datasets(
source3, target3)
self.assertEqual(actual_count3, 1)
self.assertEqual(actual_results3, (source_dataset1_1, ))
source4 = [source_dataset1_1, source_dataset1_2]
target4 = []
actual_count4, actual_results4 = DatasetAction.get_intersection_datasets(
source4, target4)
self.assertEqual(actual_count4, 0)
self.assertEqual(actual_results4, ())
source5 = []
target5 = [target_dataset1_1, target_dataset1_2]
actual_count5, actual_results5 = DatasetAction.get_intersection_datasets(
source5, target5)
self.assertEqual(actual_count5, 0)
self.assertEqual(actual_results5, ())