Skip to content

Commit 176b0ae

Browse files
tswasttswast
authored
BigQuery: add code snippets for CMEK tables (googleapis#4852)
Addition samples for BigQuery CMEK: - Create table - Update table - Query destination table - Copy table - Load table Also, removes assert for state == 'RUNNING' from snippets, as this is flaky for small jobs that complete very quickly.
1 parent 438f4a3 commit 176b0ae

File tree

2 files changed

+210
-12
lines changed

2 files changed

+210
-12
lines changed

‎docs/bigquery/snippets.py‎

Lines changed: 168 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,27 @@ def test_create_table(client, to_delete):
293293

294294
to_delete.insert(0, table)
295295

296+
def test_create_table_cmek(client, to_delete):
297+
DATASET_ID = 'create_table_cmek_{}'.format(_millis())
298+
dataset = bigquery.Dataset(client.dataset(DATASET_ID))
299+
client.create_dataset(dataset)
300+
to_delete.append(dataset)
301+
302+
# [START bigquery_create_table_cmek]
303+
table_ref = dataset.table('my_table')
304+
table = bigquery.Table(table_ref)
305+
306+
# Set the encryption key to use for the table.
307+
# TODO: Replace this key with a key you have created in Cloud KMS.
308+
kms_key_name = 'projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}'.format(
309+
'cloud-samples-tests', 'us-central1', 'test', 'test')
310+
table.encryption_configuration = bigquery.EncryptionConfiguration(
311+
kms_key_name=kms_key_name)
312+
313+
table = client.create_table(table) # API request
314+
315+
assert table.encryption_configuration.kms_key_name == kms_key_name
316+
# [END bigquery_create_table_cmek]
296317

297318
def test_get_table(client, to_delete):
298319
"""Reload a table's metadata."""
@@ -415,6 +436,40 @@ def test_update_table_multiple_properties(client, to_delete):
415436
# [END update_table_multiple_properties]
416437

417438

439+
def test_update_table_cmek(client, to_delete):
440+
"""Patch a table's metadata."""
441+
dataset_id = 'update_table_cmek_{}'.format(_millis())
442+
table_id = 'update_table_cmek_{}'.format(_millis())
443+
dataset = bigquery.Dataset(client.dataset(dataset_id))
444+
client.create_dataset(dataset)
445+
to_delete.append(dataset)
446+
447+
table = bigquery.Table(dataset.table(table_id))
448+
original_kms_key_name = (
449+
'projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}'.format(
450+
'cloud-samples-tests', 'us-central1', 'test', 'test'))
451+
table.encryption_configuration = bigquery.EncryptionConfiguration(
452+
kms_key_name=original_kms_key_name)
453+
table = client.create_table(table)
454+
to_delete.insert(0, table)
455+
456+
# [START bigquery_update_table_cmek]
457+
assert table.encryption_configuration.kms_key_name == original_kms_key_name
458+
459+
# Set a new encryption key to use for the destination.
460+
# TODO: Replace this key with a key you have created in KMS.
461+
updated_kms_key_name = 'projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}'.format(
462+
'cloud-samples-tests', 'us-central1', 'test', 'otherkey')
463+
table.encryption_configuration = bigquery.EncryptionConfiguration(
464+
kms_key_name=updated_kms_key_name)
465+
466+
table = client.update_table(table, ['encryption_configuration']) # API request
467+
468+
assert table.encryption_configuration.kms_key_name == updated_kms_key_name
469+
assert original_kms_key_name != updated_kms_key_name
470+
# [END bigquery_update_table_cmek]
471+
472+
418473
def test_table_insert_rows(client, to_delete):
419474
"""Insert / fetch table data."""
420475
dataset_id = 'table_insert_rows_dataset_{}'.format(_millis())
@@ -515,7 +570,6 @@ def test_load_table_from_uri(client, to_delete):
515570
dataset_ref.table('us_states'),
516571
job_config=job_config) # API request
517572

518-
assert load_job.state == 'RUNNING'
519573
assert load_job.job_type == 'load'
520574

521575
load_job.result() # Waits for table load to complete.
@@ -525,6 +579,42 @@ def test_load_table_from_uri(client, to_delete):
525579
# [END bigquery_load_table_gcs_json]
526580

527581

582+
def test_load_table_from_uri_cmek(client, to_delete):
583+
dataset_id = 'load_table_from_uri_cmek_{}'.format(_millis())
584+
dataset = bigquery.Dataset(client.dataset(dataset_id))
585+
client.create_dataset(dataset)
586+
to_delete.append(dataset)
587+
588+
# [START bigquery_load_table_gcs_json_cmek]
589+
# dataset_id = 'my_dataset'
590+
dataset_ref = client.dataset(dataset_id)
591+
job_config = bigquery.LoadJobConfig()
592+
job_config.autodetect = True
593+
job_config.source_format = 'NEWLINE_DELIMITED_JSON'
594+
595+
# Set the encryption key to use for the destination.
596+
# TODO: Replace this key with a key you have created in KMS.
597+
kms_key_name = 'projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}'.format(
598+
'cloud-samples-tests', 'us-central1', 'test', 'test')
599+
encryption_config = bigquery.EncryptionConfiguration(
600+
kms_key_name=kms_key_name)
601+
job_config.destination_encryption_configuration = encryption_config
602+
603+
load_job = client.load_table_from_uri(
604+
'gs://cloud-samples-data/bigquery/us-states/us-states.json',
605+
dataset_ref.table('us_states'),
606+
job_config=job_config) # API request
607+
608+
assert load_job.job_type == 'load'
609+
610+
load_job.result() # Waits for table load to complete.
611+
612+
assert load_job.state == 'DONE'
613+
table = client.get_table(dataset_ref.table('us_states'))
614+
assert table.encryption_configuration.kms_key_name == kms_key_name
615+
# [END bigquery_load_table_gcs_json_cmek]
616+
617+
528618
def test_load_table_from_uri_autodetect(client, to_delete):
529619
dataset_id = 'load_table_dataset_{}'.format(_millis())
530620
dataset = bigquery.Dataset(client.dataset(dataset_id))
@@ -543,7 +633,6 @@ def test_load_table_from_uri_autodetect(client, to_delete):
543633
dataset_ref.table('us_states'),
544634
job_config=job_config) # API request
545635

546-
assert load_job.state == 'RUNNING'
547636
assert load_job.job_type == 'load'
548637

549638
load_job.result() # Waits for table load to complete.
@@ -581,7 +670,6 @@ def test_load_table_from_uri_append(client, to_delete):
581670
table_ref,
582671
job_config=job_config) # API request
583672

584-
assert load_job.state == 'RUNNING'
585673
assert load_job.job_type == 'load'
586674

587675
load_job.result() # Waits for table load to complete.
@@ -621,7 +709,6 @@ def test_load_table_from_uri_truncate(client, to_delete):
621709
table_ref,
622710
job_config=job_config) # API request
623711

624-
assert load_job.state == 'RUNNING'
625712
assert load_job.job_type == 'load'
626713

627714
load_job.result() # Waits for table load to complete.
@@ -657,27 +744,63 @@ def _write_csv_to_storage(bucket_name, blob_name, header_row, data_rows):
657744

658745

659746
def test_copy_table(client, to_delete):
660-
DATASET_ID = 'copy_table_dataset_{}'.format(_millis())
747+
dataset_id = 'copy_table_dataset_{}'.format(_millis())
748+
dest_dataset = bigquery.Dataset(client.dataset(dataset_id))
749+
dest_dataset = client.create_dataset(dest_dataset)
750+
to_delete.append(dest_dataset)
751+
661752
# [START copy_table]
662753
source_dataset = bigquery.DatasetReference(
663754
'bigquery-public-data', 'samples')
664755
source_table_ref = source_dataset.table('shakespeare')
665756

666-
dest_dataset = bigquery.Dataset(client.dataset(DATASET_ID))
667-
dest_dataset = client.create_dataset(dest_dataset) # API request
757+
# dataset_id = 'my_dataset'
668758
dest_table_ref = dest_dataset.table('destination_table')
669759

670-
job_config = bigquery.CopyJobConfig()
671-
job = client.copy_table(
672-
source_table_ref, dest_table_ref, job_config=job_config) # API request
760+
job = client.copy_table(source_table_ref, dest_table_ref) # API request
673761
job.result() # Waits for job to complete.
674762

675763
assert job.state == 'DONE'
676764
dest_table = client.get_table(dest_table_ref) # API request
677765
assert dest_table.table_id == 'destination_table'
678766
# [END copy_table]
679767

768+
to_delete.insert(0, dest_table)
769+
770+
771+
def test_copy_table_cmek(client, to_delete):
772+
dataset_id = 'copy_table_cmek_{}'.format(_millis())
773+
dest_dataset = bigquery.Dataset(client.dataset(dataset_id))
774+
dest_dataset = client.create_dataset(dest_dataset)
680775
to_delete.append(dest_dataset)
776+
777+
# [START bigquery_copy_table_cmek]
778+
source_dataset = bigquery.DatasetReference(
779+
'bigquery-public-data', 'samples')
780+
source_table_ref = source_dataset.table('shakespeare')
781+
782+
# dataset_id = 'my_dataset'
783+
dest_dataset_ref = client.dataset(dataset_id)
784+
dest_table_ref = dest_dataset_ref.table('destination_table')
785+
786+
# Set the encryption key to use for the destination.
787+
# TODO: Replace this key with a key you have created in KMS.
788+
kms_key_name = 'projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}'.format(
789+
'cloud-samples-tests', 'us-central1', 'test', 'test')
790+
encryption_config = bigquery.EncryptionConfiguration(
791+
kms_key_name=kms_key_name)
792+
job_config = bigquery.CopyJobConfig()
793+
job_config.destination_encryption_configuration = encryption_config
794+
795+
job = client.copy_table(
796+
source_table_ref, dest_table_ref, job_config=job_config) # API request
797+
job.result() # Waits for job to complete.
798+
799+
assert job.state == 'DONE'
800+
dest_table = client.get_table(dest_table_ref)
801+
assert dest_table.encryption_configuration.kms_key_name == kms_key_name
802+
# [END bigquery_copy_table_cmek]
803+
681804
to_delete.insert(0, dest_table)
682805

683806

@@ -771,7 +894,6 @@ def test_client_query(client):
771894
'LIMIT 100')
772895
TIMEOUT = 30 # in seconds
773896
query_job = client.query(QUERY) # API request - starts the query
774-
assert query_job.state == 'RUNNING'
775897

776898
# Waits for the query to finish
777899
iterator = query_job.result(timeout=TIMEOUT)
@@ -826,6 +948,41 @@ def test_client_query_destination_table(client, to_delete):
826948
# [END bigquery_query_destination_table]
827949

828950

951+
def test_client_query_destination_table_cmek(client, to_delete):
952+
"""Run a query"""
953+
dataset_id = 'query_destination_table_{}'.format(_millis())
954+
dataset_ref = client.dataset(dataset_id)
955+
to_delete.append(dataset_ref)
956+
client.create_dataset(bigquery.Dataset(dataset_ref))
957+
to_delete.insert(0, dataset_ref.table('your_table_id'))
958+
959+
# [START bigquery_query_destination_table_cmek]
960+
job_config = bigquery.QueryJobConfig()
961+
962+
# Set the destination table. Here, dataset_id is a string, such as:
963+
# dataset_id = 'your_dataset_id'
964+
table_ref = client.dataset(dataset_id).table('your_table_id')
965+
job_config.destination = table_ref
966+
967+
# Set the encryption key to use for the destination.
968+
# TODO: Replace this key with a key you have created in KMS.
969+
kms_key_name = 'projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}'.format(
970+
'cloud-samples-tests', 'us-central1', 'test', 'test')
971+
encryption_config = bigquery.EncryptionConfiguration(
972+
kms_key_name=kms_key_name)
973+
job_config.destination_encryption_configuration = encryption_config
974+
975+
# Start the query, passing in the extra configuration.
976+
query_job = client.query(
977+
'SELECT 17 AS my_col;', job_config=job_config)
978+
query_job.result()
979+
980+
# The destination table is written using the encryption configuration.
981+
table = client.get_table(table_ref)
982+
assert table.encryption_configuration.kms_key_name == kms_key_name
983+
# [END bigquery_query_destination_table_cmek]
984+
985+
829986
def test_client_query_w_param(client):
830987
"""Run a query using a query parameter"""
831988

@@ -841,7 +998,6 @@ def test_client_query_w_param(client):
841998
job_config.query_parameters = [param]
842999
query_job = client.query(
8431000
QUERY_W_PARAM, job_config=job_config) # API request - starts the query
844-
assert query_job.state == 'RUNNING'
8451001

8461002
# Waits for the query to finish
8471003
iterator = query_job.result(timeout=TIMEOUT)

‎docs/bigquery/usage.rst‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,48 @@ Overwrite / replace an existing table with a JSON file from Cloud Storage:
235235
:start-after: [START bigquery_load_table_gcs_json_truncate]
236236
:end-before: [END bigquery_load_table_gcs_json_truncate]
237237

238+
Customer Managed Encryption Keys
239+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
240+
241+
Table data is always encrypted at rest, but BigQuery also provides a way for
242+
you to control what keys it uses to encrypt they data. See `Protecting data
243+
with Cloud KMS keys
244+
<https://cloud-dot-devsite.googleplex.com/bigquery/docs/customer-managed-encryption>`_
245+
in the BigQuery documentation for more details.
246+
247+
Create a new table, using a customer-managed encryption key from
248+
Cloud KMS to encrypt it.
249+
250+
.. literalinclude:: snippets.py
251+
:start-after: [START bigquery_create_table_cmek]
252+
:end-before: [END bigquery_create_table_cmek]
253+
254+
Change the key used to encrypt a table.
255+
256+
.. literalinclude:: snippets.py
257+
:start-after: [START bigquery_update_table_cmek]
258+
:end-before: [END bigquery_update_table_cmek]
259+
260+
Load a file from Cloud Storage, using a customer-managed encryption key from
261+
Cloud KMS for the destination table.
262+
263+
.. literalinclude:: snippets.py
264+
:start-after: [START bigquery_load_table_gcs_json_cmek]
265+
:end-before: [END bigquery_load_table_gcs_json_cmek]
266+
267+
Copy a table, using a customer-managed encryption key from Cloud KMS for the
268+
destination table.
269+
270+
.. literalinclude:: snippets.py
271+
:start-after: [START bigquery_copy_table_cmek]
272+
:end-before: [END bigquery_copy_table_cmek]
273+
274+
Write query results to a table, using a customer-managed encryption key from
275+
Cloud KMS for the destination table.
276+
277+
.. literalinclude:: snippets.py
278+
:start-after: [START bigquery_query_destination_table_cmek]
279+
:end-before: [END bigquery_query_destination_table_cmek]
238280

239281
Queries
240282
-------

0 commit comments

Comments
 (0)