feat: add samples for CMEK support (#275)

* feat: add samples for CMEK support

* test: fix backups cleanup

* test: correctly use database id for cmek restore

* test: add clean up for databases

* refactor: remove version time from sample

* refactor: use user-provided key for creating encrypted backup message

Co-authored-by: larkee <larkee@users.noreply.github.com>

samples/samples/backup_sample.py

@@ -55,6 +55,42 @@ def create_backup(instance_id, database_id, backup_id, version_time):
 
 # [END spanner_create_backup]
 
+# [START spanner_create_backup_with_encryption_key]
+def create_backup_with_encryption_key(instance_id, database_id, backup_id, kms_key_name):
+    """Creates a backup for a database using a Customer Managed Encryption Key (CMEK)."""
+    from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    # Create a backup
+    expire_time = datetime.utcnow() + timedelta(days=14)
+    encryption_config = {
+        'encryption_type': CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
+        'kms_key_name': kms_key_name,
+    }
+    backup = instance.backup(backup_id, database=database, expire_time=expire_time, encryption_config=encryption_config)
+    operation = backup.create()
+
+    # Wait for backup operation to complete.
+    operation.result(1200)
+
+    # Verify that the backup is ready.
+    backup.reload()
+    assert backup.is_ready() is True
+
+    # Get the name, create time, backup size and encryption key.
+    backup.reload()
+    print(
+        "Backup {} of size {} bytes was created at {} using encryption key {}".format(
+            backup.name, backup.size_bytes, backup.create_time, kms_key_name
+        )
+    )
+
+
+# [END spanner_create_backup_with_encryption_key]
+
 
 # [START spanner_restore_backup]
 def restore_database(instance_id, new_database_id, backup_id):
@@ -87,6 +123,42 @@ def restore_database(instance_id, new_database_id, backup_id):
 # [END spanner_restore_backup]
 
 
+# [START spanner_restore_backup_with_encryption_key]
+def restore_database_with_encryption_key(instance_id, new_database_id, backup_id, kms_key_name):
+    """Restores a database from a backup using a Customer Managed Encryption Key (CMEK)."""
+    from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+
+    # Start restoring an existing backup to a new database.
+    backup = instance.backup(backup_id)
+    encryption_config = {
+        'encryption_type': RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
+        'kms_key_name': kms_key_name,
+    }
+    new_database = instance.database(new_database_id, encryption_config=encryption_config)
+    operation = new_database.restore(backup)
+
+    # Wait for restore operation to complete.
+    operation.result(1600)
+
+    # Newly created database has restore information.
+    new_database.reload()
+    restore_info = new_database.restore_info
+    print(
+        "Database {} restored to {} from backup {} with using encryption key {}.".format(
+            restore_info.backup_info.source_database,
+            new_database_id,
+            restore_info.backup_info.backup,
+            new_database.encryption_config.kms_key_name,
+        )
+    )
+
+
+# [END spanner_restore_backup_with_encryption_key]
+
+
 # [START spanner_cancel_backup_create]
 def cancel_backup(instance_id, database_id, backup_id):
     spanner_client = spanner.Client()

samples/samples/backup_sample_test.py

@@ -38,9 +38,11 @@ def unique_backup_id():
 
 INSTANCE_ID = unique_instance_id()
 DATABASE_ID = unique_database_id()
-RETENTION_DATABASE_ID = unique_database_id()
 RESTORE_DB_ID = unique_database_id()
 BACKUP_ID = unique_backup_id()
+CMEK_RESTORE_DB_ID = unique_database_id()
+CMEK_BACKUP_ID = unique_backup_id()
+RETENTION_DATABASE_ID = unique_database_id()
 RETENTION_PERIOD = "7d"
 
 
@@ -54,6 +56,12 @@ def spanner_instance():
     op = instance.create()
     op.result(120)  # block until completion
     yield instance
+    for database_pb in instance.list_databases():
+        database = instance.database(database_pb.name.split("/")[-1])
+        database.drop()
+    for backup_pb in instance.list_backups():
+        backup = instance.backup(backup_pb.name.split("/")[-1])
+        backup.delete()
     instance.delete()
 
 
@@ -77,6 +85,16 @@ def test_create_backup(capsys, database):
     assert BACKUP_ID in out
 
 
+def test_create_backup_with_encryption_key(capsys, spanner_instance, database):
+    kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
+        spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-cmek"
+    )
+    backup_sample.create_backup_with_encryption_key(INSTANCE_ID, DATABASE_ID, CMEK_BACKUP_ID, kms_key_name)
+    out, _ = capsys.readouterr()
+    assert CMEK_BACKUP_ID in out
+    assert kms_key_name in out
+
+
 # Depends on test_create_backup having run first
 @RetryErrors(exception=DeadlineExceeded, max_tries=2)
 def test_restore_database(capsys):
@@ -87,6 +105,20 @@ def test_restore_database(capsys):
     assert BACKUP_ID in out
 
 
+# Depends on test_create_backup having run first
+@RetryErrors(exception=DeadlineExceeded, max_tries=2)
+def test_restore_database_with_encryption_key(capsys, spanner_instance):
+    kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
+        spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-cmek"
+    )
+    backup_sample.restore_database_with_encryption_key(INSTANCE_ID, CMEK_RESTORE_DB_ID, CMEK_BACKUP_ID, kms_key_name)
+    out, _ = capsys.readouterr()
+    assert (DATABASE_ID + " restored to ") in out
+    assert (CMEK_RESTORE_DB_ID + " from backup ") in out
+    assert CMEK_BACKUP_ID in out
+    assert kms_key_name in out
+
+
 # Depends on test_create_backup having run first
 def test_list_backup_operations(capsys, spanner_instance):
     backup_sample.list_backup_operations(INSTANCE_ID, DATABASE_ID)

samples/samples/snippets.py

@@ -92,6 +92,43 @@ def create_database(instance_id, database_id):
 # [END spanner_create_database]
 
 
+# [START spanner_create_database_with_encryption_key]
+def create_database_with_encryption_key(instance_id, database_id, kms_key_name):
+    """Creates a database with tables using a Customer Managed Encryption Key (CMEK)."""
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+
+    database = instance.database(
+        database_id,
+        ddl_statements=[
+            """CREATE TABLE Singers (
+            SingerId     INT64 NOT NULL,
+            FirstName    STRING(1024),
+            LastName     STRING(1024),
+            SingerInfo   BYTES(MAX)
+        ) PRIMARY KEY (SingerId)""",
+            """CREATE TABLE Albums (
+            SingerId     INT64 NOT NULL,
+            AlbumId      INT64 NOT NULL,
+            AlbumTitle   STRING(MAX)
+        ) PRIMARY KEY (SingerId, AlbumId),
+        INTERLEAVE IN PARENT Singers ON DELETE CASCADE""",
+        ],
+        encryption_config={'kms_key_name': kms_key_name},
+    )
+
+    operation = database.create()
+
+    print("Waiting for operation to complete...")
+    operation.result(120)
+
+    print("Database {} created with encryption key {}".format(
+        database.name, database.encryption_config.kms_key_name))
+
+
+# [END spanner_create_database_with_encryption_key]
+
+
 # [START spanner_insert_data]
 def insert_data(instance_id, database_id):
     """Inserts sample data into the given database.

samples/samples/snippets_test.py

@@ -33,6 +33,7 @@ def unique_database_id():
 
 INSTANCE_ID = unique_instance_id()
 DATABASE_ID = unique_database_id()
+CMEK_DATABASE_ID = unique_database_id()
 
 
 @pytest.fixture(scope="module")
@@ -63,6 +64,16 @@ def test_create_database(database):
     database.reload()
 
 
+def test_create_database_with_encryption_config(capsys, spanner_instance):
+    kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
+        spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-cmek"
+    )
+    snippets.create_database_with_encryption_key(INSTANCE_ID, CMEK_DATABASE_ID, kms_key_name)
+    out, _ = capsys.readouterr()
+    assert CMEK_DATABASE_ID in out
+    assert kms_key_name in out
+
+
 def test_insert_data(capsys):
     snippets.insert_data(INSTANCE_ID, DATABASE_ID)
     out, _ = capsys.readouterr()