Skip to content
This repository was archived by the owner on Mar 31, 2026. It is now read-only.

Commit 24fa244

Browse files
asthamohtaasthamohta
authored
feat: fgac support and samples (#867)
* feat:fgac changes and samples * linting * fixing samples * linting * linting * Update database.py * Update pool.py * Update snippets.py
1 parent 8128a87 commit 24fa244

File tree

13 files changed

+738
-35
lines changed

13 files changed

+738
-35
lines changed

‎google/cloud/spanner_dbapi/parse_utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@
151151

152152
# DDL statements follow
153153
# https://cloud.google.com/spanner/docs/data-definition-language
154-
RE_DDL = re.compile(r"^\s*(CREATE|ALTER|DROP)", re.IGNORECASE | re.DOTALL)
154+
RE_DDL = re.compile(r"^\s*(CREATE|ALTER|DROP|GRANT|REVOKE)", re.IGNORECASE | re.DOTALL)
155155

156156
RE_IS_INSERT = re.compile(r"^\s*(INSERT)", re.IGNORECASE | re.DOTALL)
157157

‎google/cloud/spanner_v1/database.py‎

Lines changed: 95 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,12 @@
2727
from google.cloud.exceptions import NotFound
2828
from google.api_core.exceptions import Aborted
2929
from google.api_core import gapic_v1
30+
from google.iam.v1 import iam_policy_pb2
31+
from google.iam.v1 import options_pb2
3032

3133
from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest
3234
from google.cloud.spanner_admin_database_v1 import Database as DatabasePB
35+
from google.cloud.spanner_admin_database_v1 import ListDatabaseRolesRequest
3336
from google.cloud.spanner_admin_database_v1 import EncryptionConfig
3437
from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig
3538
from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest
@@ -119,7 +122,8 @@ class Database(object):
119122
:class:`~google.cloud.spanner_admin_database_v1.types.DatabaseDialect`
120123
:param database_dialect:
121124
(Optional) database dialect for the database
122-
125+
:type database_role: str or None
126+
:param database_role: (Optional) user-assigned database_role for the session.
123127
"""
124128

125129
_spanner_api = None
@@ -133,6 +137,7 @@ def __init__(
133137
logger=None,
134138
encryption_config=None,
135139
database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
140+
database_role=None,
136141
):
137142
self.database_id = database_id
138143
self._instance = instance
@@ -149,9 +154,10 @@ def __init__(
149154
self._logger = logger
150155
self._encryption_config = encryption_config
151156
self._database_dialect = database_dialect
157+
self._database_role = database_role
152158

153159
if pool is None:
154-
pool = BurstyPool()
160+
pool = BurstyPool(database_role=database_role)
155161

156162
self._pool = pool
157163
pool.bind(self)
@@ -314,6 +320,14 @@ def database_dialect(self):
314320
"""
315321
return self._database_dialect
316322

323+
@property
324+
def database_role(self):
325+
"""User-assigned database_role for sessions created by the pool.
326+
:rtype: str
327+
:returns: a str with the name of the database role.
328+
"""
329+
return self._database_role
330+
317331
@property
318332
def logger(self):
319333
"""Logger used by the database.
@@ -584,16 +598,22 @@ def execute_pdml():
584598

585599
return _retry_on_aborted(execute_pdml, DEFAULT_RETRY_BACKOFF)()
586600

587-
def session(self, labels=None):
601+
def session(self, labels=None, database_role=None):
588602
"""Factory to create a session for this database.
589603
590604
:type labels: dict (str -> str) or None
591605
:param labels: (Optional) user-assigned labels for the session.
592606
607+
:type database_role: str
608+
:param database_role: (Optional) user-assigned database_role for the session.
609+
593610
:rtype: :class:`~google.cloud.spanner_v1.session.Session`
594611
:returns: a session bound to this database.
595612
"""
596-
return Session(self, labels=labels)
613+
# If role is specified in param, then that role is used
614+
# instead.
615+
role = database_role or self._database_role
616+
return Session(self, labels=labels, database_role=role)
597617

598618
def snapshot(self, **kw):
599619
"""Return an object which wraps a snapshot.
@@ -772,6 +792,29 @@ def list_database_operations(self, filter_="", page_size=None):
772792
filter_=database_filter, page_size=page_size
773793
)
774794

795+
def list_database_roles(self, page_size=None):
796+
"""Lists Cloud Spanner database roles.
797+
798+
:type page_size: int
799+
:param page_size:
800+
Optional. The maximum number of database roles in each page of results
801+
from this request. Non-positive values are ignored. Defaults to a
802+
sensible value set by the API.
803+
804+
:type: Iterable
805+
:returns:
806+
Iterable of :class:`~google.cloud.spanner_admin_database_v1.types.spanner_database_admin.DatabaseRole`
807+
resources within the current database.
808+
"""
809+
api = self._instance._client.database_admin_api
810+
metadata = _metadata_with_prefix(self.name)
811+
812+
request = ListDatabaseRolesRequest(
813+
parent=self.name,
814+
page_size=page_size,
815+
)
816+
return api.list_database_roles(request=request, metadata=metadata)
817+
775818
def table(self, table_id):
776819
"""Factory to create a table object within this database.
777820
@@ -811,6 +854,54 @@ def list_tables(self):
811854
for row in results:
812855
yield self.table(row[0])
813856

857+
def get_iam_policy(self, policy_version=None):
858+
"""Gets the access control policy for a database resource.
859+
860+
:type policy_version: int
861+
:param policy_version:
862+
(Optional) the maximum policy version that will be
863+
used to format the policy. Valid values are 0, 1 ,3.
864+
865+
:rtype: :class:`~google.iam.v1.policy_pb2.Policy`
866+
:returns:
867+
returns an Identity and Access Management (IAM) policy. It is used to
868+
specify access control policies for Cloud Platform
869+
resources.
870+
"""
871+
api = self._instance._client.database_admin_api
872+
metadata = _metadata_with_prefix(self.name)
873+
874+
request = iam_policy_pb2.GetIamPolicyRequest(
875+
resource=self.name,
876+
options=options_pb2.GetPolicyOptions(
877+
requested_policy_version=policy_version
878+
),
879+
)
880+
response = api.get_iam_policy(request=request, metadata=metadata)
881+
return response
882+
883+
def set_iam_policy(self, policy):
884+
"""Sets the access control policy on a database resource.
885+
Replaces any existing policy.
886+
887+
:type policy: :class:`~google.iam.v1.policy_pb2.Policy`
888+
:param policy_version:
889+
the complete policy to be applied to the resource.
890+
891+
:rtype: :class:`~google.iam.v1.policy_pb2.Policy`
892+
:returns:
893+
returns the new Identity and Access Management (IAM) policy.
894+
"""
895+
api = self._instance._client.database_admin_api
896+
metadata = _metadata_with_prefix(self.name)
897+
898+
request = iam_policy_pb2.SetIamPolicyRequest(
899+
resource=self.name,
900+
policy=policy,
901+
)
902+
response = api.set_iam_policy(request=request, metadata=metadata)
903+
return response
904+
814905

815906
class BatchCheckout(object):
816907
"""Context manager for using a batch from a database.

‎google/cloud/spanner_v1/instance.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -431,6 +431,7 @@ def database(
431431
logger=None,
432432
encryption_config=None,
433433
database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
434+
database_role=None,
434435
):
435436
"""Factory to create a database within this instance.
436437
@@ -477,6 +478,7 @@ def database(
477478
logger=logger,
478479
encryption_config=encryption_config,
479480
database_dialect=database_dialect,
481+
database_role=database_role,
480482
)
481483

482484
def list_databases(self, page_size=None):

‎google/cloud/spanner_v1/pool.py‎

Lines changed: 76 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import queue
1919

2020
from google.cloud.exceptions import NotFound
21+
from google.cloud.spanner_v1 import BatchCreateSessionsRequest
22+
from google.cloud.spanner_v1 import Session
2123
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
2224

2325

@@ -30,14 +32,18 @@ class AbstractSessionPool(object):
3032
:type labels: dict (str -> str) or None
3133
:param labels: (Optional) user-assigned labels for sessions created
3234
by the pool.
35+
36+
:type database_role: str
37+
:param database_role: (Optional) user-assigned database_role for the session.
3338
"""
3439

3540
_database = None
3641

37-
def __init__(self, labels=None):
42+
def __init__(self, labels=None, database_role=None):
3843
if labels is None:
3944
labels = {}
4045
self._labels = labels
46+
self._database_role = database_role
4147

4248
@property
4349
def labels(self):
@@ -48,6 +54,15 @@ def labels(self):
4854
"""
4955
return self._labels
5056

57+
@property
58+
def database_role(self):
59+
"""User-assigned database_role for sessions created by the pool.
60+
61+
:rtype: str
62+
:returns: database_role assigned by the user
63+
"""
64+
return self._database_role
65+
5166
def bind(self, database):
5267
"""Associate the pool with a database.
5368
@@ -104,9 +119,9 @@ def _new_session(self):
104119
:rtype: :class:`~google.cloud.spanner_v1.session.Session`
105120
:returns: new session instance.
106121
"""
107-
if self.labels:
108-
return self._database.session(labels=self.labels)
109-
return self._database.session()
122+
return self._database.session(
123+
labels=self.labels, database_role=self.database_role
124+
)
110125

111126
def session(self, **kwargs):
112127
"""Check out a session from the pool.
@@ -146,13 +161,22 @@ class FixedSizePool(AbstractSessionPool):
146161
:type labels: dict (str -> str) or None
147162
:param labels: (Optional) user-assigned labels for sessions created
148163
by the pool.
164+
165+
:type database_role: str
166+
:param database_role: (Optional) user-assigned database_role for the session.
149167
"""
150168

151169
DEFAULT_SIZE = 10
152170
DEFAULT_TIMEOUT = 10
153171

154-
def __init__(self, size=DEFAULT_SIZE, default_timeout=DEFAULT_TIMEOUT, labels=None):
155-
super(FixedSizePool, self).__init__(labels=labels)
172+
def __init__(
173+
self,
174+
size=DEFAULT_SIZE,
175+
default_timeout=DEFAULT_TIMEOUT,
176+
labels=None,
177+
database_role=None,
178+
):
179+
super(FixedSizePool, self).__init__(labels=labels, database_role=database_role)
156180
self.size = size
157181
self.default_timeout = default_timeout
158182
self._sessions = queue.LifoQueue(size)
@@ -167,9 +191,14 @@ def bind(self, database):
167191
self._database = database
168192
api = database.spanner_api
169193
metadata = _metadata_with_prefix(database.name)
194+
self._database_role = self._database_role or self._database.database_role
195+
request = BatchCreateSessionsRequest(
196+
session_template=Session(creator_role=self.database_role),
197+
)
170198

171199
while not self._sessions.full():
172200
resp = api.batch_create_sessions(
201+
request=request,
173202
database=database.name,
174203
session_count=self.size - self._sessions.qsize(),
175204
metadata=metadata,
@@ -243,10 +272,13 @@ class BurstyPool(AbstractSessionPool):
243272
:type labels: dict (str -> str) or None
244273
:param labels: (Optional) user-assigned labels for sessions created
245274
by the pool.
275+
276+
:type database_role: str
277+
:param database_role: (Optional) user-assigned database_role for the session.
246278
"""
247279

248-
def __init__(self, target_size=10, labels=None):
249-
super(BurstyPool, self).__init__(labels=labels)
280+
def __init__(self, target_size=10, labels=None, database_role=None):
281+
super(BurstyPool, self).__init__(labels=labels, database_role=database_role)
250282
self.target_size = target_size
251283
self._database = None
252284
self._sessions = queue.LifoQueue(target_size)
@@ -259,6 +291,7 @@ def bind(self, database):
259291
when needed.
260292
"""
261293
self._database = database
294+
self._database_role = self._database_role or self._database.database_role
262295

263296
def get(self):
264297
"""Check a session out from the pool.
@@ -340,10 +373,20 @@ class PingingPool(AbstractSessionPool):
340373
:type labels: dict (str -> str) or None
341374
:param labels: (Optional) user-assigned labels for sessions created
342375
by the pool.
376+
377+
:type database_role: str
378+
:param database_role: (Optional) user-assigned database_role for the session.
343379
"""
344380

345-
def __init__(self, size=10, default_timeout=10, ping_interval=3000, labels=None):
346-
super(PingingPool, self).__init__(labels=labels)
381+
def __init__(
382+
self,
383+
size=10,
384+
default_timeout=10,
385+
ping_interval=3000,
386+
labels=None,
387+
database_role=None,
388+
):
389+
super(PingingPool, self).__init__(labels=labels, database_role=database_role)
347390
self.size = size
348391
self.default_timeout = default_timeout
349392
self._delta = datetime.timedelta(seconds=ping_interval)
@@ -360,9 +403,15 @@ def bind(self, database):
360403
api = database.spanner_api
361404
metadata = _metadata_with_prefix(database.name)
362405
created_session_count = 0
406+
self._database_role = self._database_role or self._database.database_role
407+
408+
request = BatchCreateSessionsRequest(
409+
session_template=Session(creator_role=self.database_role),
410+
)
363411

364412
while created_session_count < self.size:
365413
resp = api.batch_create_sessions(
414+
request=request,
366415
database=database.name,
367416
session_count=self.size - created_session_count,
368417
metadata=metadata,
@@ -470,13 +519,27 @@ class TransactionPingingPool(PingingPool):
470519
:type labels: dict (str -> str) or None
471520
:param labels: (Optional) user-assigned labels for sessions created
472521
by the pool.
522+
523+
:type database_role: str
524+
:param database_role: (Optional) user-assigned database_role for the session.
473525
"""
474526

475-
def __init__(self, size=10, default_timeout=10, ping_interval=3000, labels=None):
527+
def __init__(
528+
self,
529+
size=10,
530+
default_timeout=10,
531+
ping_interval=3000,
532+
labels=None,
533+
database_role=None,
534+
):
476535
self._pending_sessions = queue.Queue()
477536

478537
super(TransactionPingingPool, self).__init__(
479-
size, default_timeout, ping_interval, labels=labels
538+
size,
539+
default_timeout,
540+
ping_interval,
541+
labels=labels,
542+
database_role=database_role,
480543
)
481544

482545
self.begin_pending_transactions()
@@ -489,6 +552,7 @@ def bind(self, database):
489552
when needed.
490553
"""
491554
super(TransactionPingingPool, self).bind(database)
555+
self._database_role = self._database_role or self._database.database_role
492556
self.begin_pending_transactions()
493557

494558
def put(self, session):

0 commit comments

Comments
 (0)