Skip to content

Commit 8abb5fa

Browse files
larkeelarkee
committed
feat: add support for creating databases with CMEK
1 parent b078413 commit 8abb5fa

File tree

6 files changed

+129
-4
lines changed

6 files changed

+129
-4
lines changed

‎google/cloud/spanner_v1/backup.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ def __init__(self, backup_id, instance, database="", expire_time=None):
6262
self._size_bytes = None
6363
self._state = None
6464
self._referencing_databases = None
65+
self._encryption_info = None
6566

6667
@property
6768
def name(self):
@@ -137,6 +138,14 @@ def referencing_databases(self):
137138
"""
138139
return self._referencing_databases
139140

141+
@property
142+
def encryption_info(self):
143+
"""Encryption info for this backup.
144+
:rtype: :class:`~google.clod.spanner_admin_database_v1.proto.common_pb2.EncryptionInfo`
145+
:returns: a class representing the encryption info
146+
"""
147+
return self._encryption_info
148+
140149
@classmethod
141150
def from_pb(cls, backup_pb, instance):
142151
"""Create an instance of this class from a protobuf message.
@@ -231,6 +240,7 @@ def reload(self):
231240
self._size_bytes = pb.size_bytes
232241
self._state = BackupPB.State(pb.state)
233242
self._referencing_databases = pb.referencing_databases
243+
self._encryption_info = pb.encryption_info
234244

235245
def update_expire_time(self, new_expire_time):
236246
"""Update the expire time of this backup.

‎google/cloud/spanner_v1/database.py

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
SpannerGrpcTransport,
4747
)
4848
from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest
49+
from google.cloud.spanner_admin_database_v1 import EncryptionConfig
4950
from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest
5051
from google.cloud.spanner_v1 import ExecuteSqlRequest
5152
from google.cloud.spanner_v1 import (
@@ -95,11 +96,26 @@ class Database(object):
9596
:param pool: (Optional) session pool to be used by database. If not
9697
passed, the database will construct an instance of
9798
:class:`~google.cloud.spanner_v1.pool.BurstyPool`.
99+
100+
:type encryption_config:
101+
:class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
102+
or :class:`dict`
103+
:param encryption_config:
104+
(Optional) Encryption information about the database.
105+
If a dict is provided, it must be of the same form as the protobuf
106+
message :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
98107
"""
99108

100109
_spanner_api = None
101110

102-
def __init__(self, database_id, instance, ddl_statements=(), pool=None):
111+
def __init__(
112+
self,
113+
database_id,
114+
instance,
115+
ddl_statements=(),
116+
pool=None,
117+
encryption_config=None,
118+
):
103119
self.database_id = database_id
104120
self._instance = instance
105121
self._ddl_statements = _check_ddl_statements(ddl_statements)
@@ -108,6 +124,13 @@ def __init__(self, database_id, instance, ddl_statements=(), pool=None):
108124
self._create_time = None
109125
self._restore_info = None
110126

127+
if type(encryption_config) == dict:
128+
self._encryption_config = EncryptionConfig(
129+
kms_key_name=encryption_config["kms_key_name"]
130+
)
131+
else:
132+
self._encryption_config = encryption_config
133+
111134
if pool is None:
112135
pool = BurstyPool()
113136

@@ -204,6 +227,14 @@ def restore_info(self):
204227
"""
205228
return self._restore_info
206229

230+
@property
231+
def encryption_config(self):
232+
"""Encryption config for this database.
233+
:rtype: :class:`~google.cloud.spanner_admin_instance_v1.proto.common_pb2.EncryptionConfig`
234+
:returns: an object representing the restore info for this database
235+
"""
236+
return self._encryption_config
237+
207238
@property
208239
def ddl_statements(self):
209240
"""DDL Statements used to define database schema.
@@ -273,6 +304,7 @@ def create(self):
273304
parent=self._instance.name,
274305
create_statement="CREATE DATABASE %s" % (db_name,),
275306
extra_statements=list(self._ddl_statements),
307+
encryption_config=self._encryption_config,
276308
)
277309
future = api.create_database(request=request, metadata=metadata)
278310
return future
@@ -313,6 +345,7 @@ def reload(self):
313345
self._state = DatabasePB.State(response.state)
314346
self._create_time = response.create_time
315347
self._restore_info = response.restore_info
348+
self._encryption_config = response.encryption_config
316349

317350
def update_ddl(self, ddl_statements, operation_id=""):
318351
"""Update DDL for this database.

‎google/cloud/spanner_v1/instance.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -356,7 +356,9 @@ def delete(self):
356356

357357
api.delete_instance(name=self.name, metadata=metadata)
358358

359-
def database(self, database_id, ddl_statements=(), pool=None):
359+
def database(
360+
self, database_id, ddl_statements=(), pool=None, encryption_config=None
361+
):
360362
"""Factory to create a database within this instance.
361363
362364
:type database_id: str
@@ -370,10 +372,24 @@ def database(self, database_id, ddl_statements=(), pool=None):
370372
:class:`~google.cloud.spanner_v1.pool.AbstractSessionPool`.
371373
:param pool: (Optional) session pool to be used by database.
372374
375+
:type encryption_config:
376+
:class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
377+
or :class:`dict`
378+
:param encryption_config:
379+
(Optional) Encryption information about the database.
380+
If a dict is provided, it must be of the same form as the protobuf
381+
message :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig
382+
373383
:rtype: :class:`~google.cloud.spanner_v1.database.Database`
374384
:returns: a database owned by this instance.
375385
"""
376-
return Database(database_id, self, ddl_statements=ddl_statements, pool=pool)
386+
return Database(
387+
database_id,
388+
self,
389+
ddl_statements=ddl_statements,
390+
pool=pool,
391+
encryption_config=encryption_config,
392+
)
377393

378394
def list_databases(self, page_size=None):
379395
"""List databases for the instance.

‎tests/unit/test_backup.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,16 @@ def test_referencing_databases_property(self):
170170
expected = backup._referencing_databases = [self.DATABASE_NAME]
171171
self.assertEqual(backup.referencing_databases, expected)
172172

173+
def test_encrpytion_info_property(self):
174+
from google.cloud.spanner_admin_database_v1 import EncryptionInfo
175+
176+
instance = _Instance(self.INSTANCE_NAME)
177+
backup = self._make_one(self.BACKUP_ID, instance)
178+
expected = backup._encryption_info = EncryptionInfo(
179+
kms_key_version="kms_key_version"
180+
)
181+
self.assertEqual(backup.encryption_info, expected)
182+
173183
def test_create_grpc_error(self):
174184
from google.api_core.exceptions import GoogleAPICallError
175185
from google.api_core.exceptions import Unknown
@@ -429,8 +439,10 @@ def test_reload_not_found(self):
429439

430440
def test_reload_success(self):
431441
from google.cloud.spanner_admin_database_v1 import Backup
442+
from google.cloud.spanner_admin_database_v1 import EncryptionInfo
432443

433444
timestamp = self._make_timestamp()
445+
encryption_info = EncryptionInfo(kms_key_version="kms_key_version")
434446

435447
client = _Client()
436448
backup_pb = Backup(
@@ -441,6 +453,7 @@ def test_reload_success(self):
441453
size_bytes=10,
442454
state=1,
443455
referencing_databases=[],
456+
encryption_info=encryption_info,
444457
)
445458
api = client.database_admin_api = self._make_database_admin_api()
446459
api.get_backup.return_value = backup_pb
@@ -455,6 +468,7 @@ def test_reload_success(self):
455468
self.assertEqual(backup.size_bytes, 10)
456469
self.assertEqual(backup.state, Backup.State.CREATING)
457470
self.assertEqual(backup.referencing_databases, [])
471+
self.assertEqual(backup.encryption_info, encryption_info)
458472

459473
api.get_backup.assert_called_once_with(
460474
name=self.BACKUP_NAME,

‎tests/unit/test_database.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,32 @@ def test_ctor_w_ddl_statements_ok(self):
145145
self.assertIs(database._instance, instance)
146146
self.assertEqual(list(database.ddl_statements), DDL_STATEMENTS)
147147

148+
def test_ctor_w_encryption_config(self):
149+
from google.cloud.spanner_admin_database_v1 import EncryptionConfig
150+
151+
instance = _Instance(self.INSTANCE_NAME)
152+
encryption_config = EncryptionConfig(kms_key_name="kms_key")
153+
database = self._make_one(
154+
self.DATABASE_ID, instance, encryption_config=encryption_config
155+
)
156+
self.assertEqual(database.database_id, self.DATABASE_ID)
157+
self.assertIs(database._instance, instance)
158+
self.assertEqual(database._encryption_config, encryption_config)
159+
160+
def test_ctor_w_encryption_config_dict(self):
161+
from google.cloud.spanner_admin_database_v1 import EncryptionConfig
162+
163+
instance = _Instance(self.INSTANCE_NAME)
164+
encryption_config_dict = {"kms_key_name": "kms_key"}
165+
encryption_config = EncryptionConfig(kms_key_name="kms_key")
166+
database = self._make_one(
167+
self.DATABASE_ID, instance, encryption_config=encryption_config_dict
168+
)
169+
self.assertEqual(database.database_id, self.DATABASE_ID)
170+
self.assertIs(database._instance, instance)
171+
self.assertEqual(database._encryption_config, encryption_config)
172+
173+
148174
def test_from_pb_bad_database_name(self):
149175
from google.cloud.spanner_admin_database_v1 import Database
150176

@@ -249,6 +275,17 @@ def test_restore_info(self):
249275
)
250276
self.assertEqual(database.restore_info, restore_info)
251277

278+
def test_encryption_config(self):
279+
from google.cloud.spanner_admin_database_v1 import EncryptionConfig
280+
281+
instance = _Instance(self.INSTANCE_NAME)
282+
pool = _Pool()
283+
database = self._make_one(self.DATABASE_ID, instance, pool=pool)
284+
encryption_config = database._encryption_config = mock.create_autospec(
285+
EncryptionConfig, instance=True
286+
)
287+
self.assertEqual(database.encryption_config, encryption_config)
288+
252289
def test_spanner_api_property_w_scopeless_creds(self):
253290

254291
client = _Client()
@@ -386,6 +423,7 @@ def test_create_grpc_error(self):
386423
parent=self.INSTANCE_NAME,
387424
create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),
388425
extra_statements=[],
426+
encryption_config=None,
389427
)
390428

391429
api.create_database.assert_called_once_with(
@@ -412,6 +450,7 @@ def test_create_already_exists(self):
412450
parent=self.INSTANCE_NAME,
413451
create_statement="CREATE DATABASE `{}`".format(DATABASE_ID_HYPHEN),
414452
extra_statements=[],
453+
encryption_config=None,
415454
)
416455

417456
api.create_database.assert_called_once_with(
@@ -437,6 +476,7 @@ def test_create_instance_not_found(self):
437476
parent=self.INSTANCE_NAME,
438477
create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),
439478
extra_statements=[],
479+
encryption_config=None,
440480
)
441481

442482
api.create_database.assert_called_once_with(
@@ -466,6 +506,7 @@ def test_create_success(self):
466506
parent=self.INSTANCE_NAME,
467507
create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),
468508
extra_statements=DDL_STATEMENTS,
509+
encryption_config=None,
469510
)
470511

471512
api.create_database.assert_called_once_with(
@@ -565,6 +606,7 @@ def test_reload_not_found(self):
565606

566607
def test_reload_success(self):
567608
from google.cloud.spanner_admin_database_v1 import Database
609+
from google.cloud.spanner_admin_database_v1 import EncryptionConfig
568610
from google.cloud.spanner_admin_database_v1 import GetDatabaseDdlResponse
569611
from google.cloud.spanner_admin_database_v1 import RestoreInfo
570612
from google.cloud._helpers import _datetime_to_pb_timestamp
@@ -575,12 +617,14 @@ def test_reload_success(self):
575617

576618
client = _Client()
577619
ddl_pb = GetDatabaseDdlResponse(statements=DDL_STATEMENTS)
620+
encryption_config = EncryptionConfig(kms_key_name="kms_key")
578621
api = client.database_admin_api = self._make_database_admin_api()
579622
api.get_database_ddl.return_value = ddl_pb
580623
db_pb = Database(
581624
state=2,
582625
create_time=_datetime_to_pb_timestamp(timestamp),
583626
restore_info=restore_info,
627+
encryption_config=encryption_config,
584628
)
585629
api.get_database.return_value = db_pb
586630
instance = _Instance(self.INSTANCE_NAME, client=client)
@@ -592,6 +636,7 @@ def test_reload_success(self):
592636
self.assertEqual(database._create_time, timestamp)
593637
self.assertEqual(database._restore_info, restore_info)
594638
self.assertEqual(database._ddl_statements, tuple(DDL_STATEMENTS))
639+
self.assertEqual(database._encryption_config, encryption_config)
595640

596641
api.get_database_ddl.assert_called_once_with(
597642
database=self.DATABASE_NAME,
@@ -1290,6 +1335,7 @@ def test_context_mgr_success(self):
12901335

12911336
expected_txn_options = TransactionOptions(read_write={})
12921337

1338+
12931339
api.commit.assert_called_once_with(
12941340
session=self.SESSION_NAME,
12951341
mutations=[],

‎tests/unit/test_instance.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -488,16 +488,21 @@ def test_database_factory_defaults(self):
488488
self.assertIs(pool._database, database)
489489

490490
def test_database_factory_explicit(self):
491+
from google.cloud.spanner_admin_database_v1 import EncryptionConfig
491492
from google.cloud.spanner_v1.database import Database
492493
from tests._fixtures import DDL_STATEMENTS
493494

494495
client = _Client(self.PROJECT)
495496
instance = self._make_one(self.INSTANCE_ID, client, self.CONFIG_NAME)
496497
DATABASE_ID = "database-id"
497498
pool = _Pool()
499+
encryption_config = EncryptionConfig(kms_key_name="kms_key")
498500

499501
database = instance.database(
500-
DATABASE_ID, ddl_statements=DDL_STATEMENTS, pool=pool
502+
DATABASE_ID,
503+
ddl_statements=DDL_STATEMENTS,
504+
pool=pool,
505+
encryption_config=encryption_config,
501506
)
502507

503508
self.assertIsInstance(database, Database)
@@ -506,6 +511,7 @@ def test_database_factory_explicit(self):
506511
self.assertEqual(list(database.ddl_statements), DDL_STATEMENTS)
507512
self.assertIs(database._pool, pool)
508513
self.assertIs(pool._bound, database)
514+
self.assertIs(database._encryption_config, encryption_config)
509515

510516
def test_list_databases(self):
511517
from google.cloud.spanner_admin_database_v1 import Database as DatabasePB

0 commit comments

Comments
 (0)