googleapis · asthamohta · Dec 13, 2022 · Dec 12, 2022 · Dec 13, 2022 · Dec 13, 2022
@@ -151,7 +151,7 @@
 
 # DDL statements follow
 # https://cloud.google.com/spanner/docs/data-definition-language
-RE_DDL = re.compile(r"^\s*(CREATE|ALTER|DROP)", re.IGNORECASE | re.DOTALL)
+RE_DDL = re.compile(r"^\s*(CREATE|ALTER|DROP|GRANT|REVOKE)", re.IGNORECASE | re.DOTALL)
 
 RE_IS_INSERT = re.compile(r"^\s*(INSERT)", re.IGNORECASE | re.DOTALL)
 

@@ -27,9 +27,12 @@
 from google.cloud.exceptions import NotFound
 from google.api_core.exceptions import Aborted
 from google.api_core import gapic_v1
+from google.iam.v1 import iam_policy_pb2
+from google.iam.v1 import options_pb2
 
 from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest
 from google.cloud.spanner_admin_database_v1 import Database as DatabasePB
+from google.cloud.spanner_admin_database_v1 import ListDatabaseRolesRequest
 from google.cloud.spanner_admin_database_v1 import EncryptionConfig
 from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig
 from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest
@@ -119,7 +122,8 @@ class Database(object):
         :class:`~google.cloud.spanner_admin_database_v1.types.DatabaseDialect`
     :param database_dialect:
         (Optional) database dialect for the database
-
+    :type database_role: str or None
+    :param database_role: (Optional) user-assigned database_role for the session.
     """
 
     _spanner_api = None
@@ -133,6 +137,7 @@ def __init__(
         logger=None,
         encryption_config=None,
         database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
+        database_role=None,
     ):
         self.database_id = database_id
         self._instance = instance
@@ -149,9 +154,10 @@ def __init__(
         self._logger = logger
         self._encryption_config = encryption_config
         self._database_dialect = database_dialect
+        self._database_role = database_role
 
         if pool is None:
-            pool = BurstyPool()
+            pool = BurstyPool(database_role=database_role)
 
         self._pool = pool
         pool.bind(self)
@@ -314,6 +320,14 @@ def database_dialect(self):
         """
         return self._database_dialect
 
+    @property
+    def database_role(self):
+        """User-assigned database_role for sessions created by the pool.
+        :rtype: str
+        :returns: a str with the name of the database role.
+        """
+        return self._database_role
+
     @property
     def logger(self):
         """Logger used by the database.
@@ -584,16 +598,22 @@ def execute_pdml():
 
         return _retry_on_aborted(execute_pdml, DEFAULT_RETRY_BACKOFF)()
 
-    def session(self, labels=None):
+    def session(self, labels=None, database_role=None):
         """Factory to create a session for this database.
 
         :type labels: dict (str -> str) or None
         :param labels: (Optional) user-assigned labels for the session.
 
+        :type database_role: str
+        :param database_role: (Optional) user-assigned database_role for the session.
+
         :rtype: :class:`~google.cloud.spanner_v1.session.Session`
         :returns: a session bound to this database.
         """
-        return Session(self, labels=labels)
+        # If role is specified in param, then that role is used
+        # instead.
+        role = database_role or self._database_role
+        return Session(self, labels=labels, database_role=role)
 
     def snapshot(self, **kw):
         """Return an object which wraps a snapshot.
@@ -772,6 +792,29 @@ def list_database_operations(self, filter_="", page_size=None):
             filter_=database_filter, page_size=page_size
         )
 
+    def list_database_roles(self, page_size=None):
+        """Lists Cloud Spanner database roles.
+
+        :type page_size: int
+        :param page_size:
+            Optional. The maximum number of database roles in each page of results
+            from this request. Non-positive values are ignored. Defaults to a
+            sensible value set by the API.
+
+        :type: Iterable
+        :returns:
+            Iterable of :class:`~google.cloud.spanner_admin_database_v1.types.spanner_database_admin.DatabaseRole`
+            resources within the current database.
+        """
+        api = self._instance._client.database_admin_api
+        metadata = _metadata_with_prefix(self.name)
+
+        request = ListDatabaseRolesRequest(
+            parent=self.name,
+            page_size=page_size,
+        )
+        return api.list_database_roles(request=request, metadata=metadata)
+
     def table(self, table_id):
         """Factory to create a table object within this database.
 
@@ -811,6 +854,54 @@ def list_tables(self):
             for row in results:
                 yield self.table(row[0])
 
+    def get_iam_policy(self, policy_version=None):
+        """Gets the access control policy for a database resource.
+
+        :type policy_version: int
+        :param policy_version:
+            (Optional) the maximum policy version that will be
+            used to format the policy. Valid values are 0, 1 ,3.
+
+        :rtype: :class:`~google.iam.v1.policy_pb2.Policy`
+        :returns:
+            returns an Identity and Access Management (IAM) policy. It is used to
+            specify access control policies for Cloud Platform
+            resources.
+        """
+        api = self._instance._client.database_admin_api
+        metadata = _metadata_with_prefix(self.name)
+
+        request = iam_policy_pb2.GetIamPolicyRequest(
+            resource=self.name,
+            options=options_pb2.GetPolicyOptions(
+                requested_policy_version=policy_version
+            ),
+        )
+        response = api.get_iam_policy(request=request, metadata=metadata)
+        return response
+
+    def set_iam_policy(self, policy):
+        """Sets the access control policy on a database resource.
+        Replaces any existing policy.
+
+        :type policy: :class:`~google.iam.v1.policy_pb2.Policy`
+        :param policy_version:
+            the complete policy to be applied to the resource.
+
+        :rtype: :class:`~google.iam.v1.policy_pb2.Policy`
+        :returns:
+            returns the new Identity and Access Management (IAM) policy.
+        """
+        api = self._instance._client.database_admin_api
+        metadata = _metadata_with_prefix(self.name)
+
+        request = iam_policy_pb2.SetIamPolicyRequest(
+            resource=self.name,
+            policy=policy,
+        )
+        response = api.set_iam_policy(request=request, metadata=metadata)
+        return response
+
 
 class BatchCheckout(object):
     """Context manager for using a batch from a database.

@@ -431,6 +431,7 @@ def database(
         logger=None,
         encryption_config=None,
         database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
+        database_role=None,
     ):
         """Factory to create a database within this instance.
 
@@ -477,6 +478,7 @@ def database(
             logger=logger,
             encryption_config=encryption_config,
             database_dialect=database_dialect,
+            database_role=database_role,
         )
 
     def list_databases(self, page_size=None):

@@ -18,6 +18,8 @@
 import queue
 
 from google.cloud.exceptions import NotFound
+from google.cloud.spanner_v1 import BatchCreateSessionsRequest
+from google.cloud.spanner_v1 import Session
 from google.cloud.spanner_v1._helpers import _metadata_with_prefix
 
 
@@ -30,14 +32,18 @@ class AbstractSessionPool(object):
     :type labels: dict (str -> str) or None
     :param labels: (Optional) user-assigned labels for sessions created
                     by the pool.
+
+    :type database_role: str
+    :param database_role: (Optional) user-assigned database_role for the session.
     """
 
     _database = None
 
-    def __init__(self, labels=None):
+    def __init__(self, labels=None, database_role=None):
         if labels is None:
             labels = {}
         self._labels = labels
+        self._database_role = database_role
 
     @property
     def labels(self):
@@ -48,6 +54,15 @@ def labels(self):
         """
         return self._labels
 
+    @property
+    def database_role(self):
+        """User-assigned database_role for sessions created by the pool.
+
+        :rtype: str
+        :returns: database_role assigned by the user
+        """
+        return self._database_role
+
     def bind(self, database):
         """Associate the pool with a database.
 
@@ -104,9 +119,9 @@ def _new_session(self):
         :rtype: :class:`~google.cloud.spanner_v1.session.Session`
         :returns: new session instance.
         """
-        if self.labels:
-            return self._database.session(labels=self.labels)
-        return self._database.session()
+        return self._database.session(
+            labels=self.labels, database_role=self.database_role
+        )
 
     def session(self, **kwargs):
         """Check out a session from the pool.
@@ -146,13 +161,22 @@ class FixedSizePool(AbstractSessionPool):
     :type labels: dict (str -> str) or None
     :param labels: (Optional) user-assigned labels for sessions created
                     by the pool.
+
+    :type database_role: str
+    :param database_role: (Optional) user-assigned database_role for the session.
     """
 
     DEFAULT_SIZE = 10
     DEFAULT_TIMEOUT = 10
 
-    def __init__(self, size=DEFAULT_SIZE, default_timeout=DEFAULT_TIMEOUT, labels=None):
-        super(FixedSizePool, self).__init__(labels=labels)
+    def __init__(
+        self,
+        size=DEFAULT_SIZE,
+        default_timeout=DEFAULT_TIMEOUT,
+        labels=None,
+        database_role=None,
+    ):
+        super(FixedSizePool, self).__init__(labels=labels, database_role=database_role)
         self.size = size
         self.default_timeout = default_timeout
         self._sessions = queue.LifoQueue(size)
@@ -167,9 +191,14 @@ def bind(self, database):
         self._database = database
         api = database.spanner_api
         metadata = _metadata_with_prefix(database.name)
+        self._database_role = self._database_role or self._database.database_role
+        request = BatchCreateSessionsRequest(
+            session_template=Session(creator_role=self.database_role),
+        )
 
         while not self._sessions.full():
             resp = api.batch_create_sessions(
+                request=request,
                 database=database.name,
                 session_count=self.size - self._sessions.qsize(),
                 metadata=metadata,
@@ -243,10 +272,13 @@ class BurstyPool(AbstractSessionPool):
     :type labels: dict (str -> str) or None
     :param labels: (Optional) user-assigned labels for sessions created
                     by the pool.
+
+    :type database_role: str
+    :param database_role: (Optional) user-assigned database_role for the session.
     """
 
-    def __init__(self, target_size=10, labels=None):
-        super(BurstyPool, self).__init__(labels=labels)
+    def __init__(self, target_size=10, labels=None, database_role=None):
+        super(BurstyPool, self).__init__(labels=labels, database_role=database_role)
         self.target_size = target_size
         self._database = None
         self._sessions = queue.LifoQueue(target_size)
@@ -259,6 +291,7 @@ def bind(self, database):
                          when needed.
         """
         self._database = database
+        self._database_role = self._database_role or self._database.database_role
 
     def get(self):
         """Check a session out from the pool.
@@ -340,10 +373,20 @@ class PingingPool(AbstractSessionPool):
     :type labels: dict (str -> str) or None
     :param labels: (Optional) user-assigned labels for sessions created
                     by the pool.
+
+    :type database_role: str
+    :param database_role: (Optional) user-assigned database_role for the session.
     """
 
-    def __init__(self, size=10, default_timeout=10, ping_interval=3000, labels=None):
-        super(PingingPool, self).__init__(labels=labels)
+    def __init__(
+        self,
+        size=10,
+        default_timeout=10,
+        ping_interval=3000,
+        labels=None,
+        database_role=None,
+    ):
+        super(PingingPool, self).__init__(labels=labels, database_role=database_role)
         self.size = size
         self.default_timeout = default_timeout
         self._delta = datetime.timedelta(seconds=ping_interval)
@@ -360,9 +403,15 @@ def bind(self, database):
         api = database.spanner_api
         metadata = _metadata_with_prefix(database.name)
         created_session_count = 0
+        self._database_role = self._database_role or self._database.database_role
+
+        request = BatchCreateSessionsRequest(
+            session_template=Session(creator_role=self.database_role),
+        )
 
         while created_session_count < self.size:
             resp = api.batch_create_sessions(
+                request=request,
                 database=database.name,
                 session_count=self.size - created_session_count,
                 metadata=metadata,
@@ -470,13 +519,27 @@ class TransactionPingingPool(PingingPool):
     :type labels: dict (str -> str) or None
     :param labels: (Optional) user-assigned labels for sessions created
                     by the pool.
+
+    :type database_role: str
+    :param database_role: (Optional) user-assigned database_role for the session.
     """
 
-    def __init__(self, size=10, default_timeout=10, ping_interval=3000, labels=None):
+    def __init__(
+        self,
+        size=10,
+        default_timeout=10,
+        ping_interval=3000,
+        labels=None,
+        database_role=None,
+    ):
         self._pending_sessions = queue.Queue()
 
         super(TransactionPingingPool, self).__init__(
-            size, default_timeout, ping_interval, labels=labels
+            size,
+            default_timeout,
+            ping_interval,
+            labels=labels,
+            database_role=database_role,
         )
 
         self.begin_pending_transactions()
@@ -489,6 +552,7 @@ def bind(self, database):
                          when needed.
         """
         super(TransactionPingingPool, self).bind(database)
+        self._database_role = self._database_role or self._database.database_role
         self.begin_pending_transactions()
 
     def put(self, session):