feat: add configurable leader placement support (#399)

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-spanner/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

Author: zoercai

google/cloud/spanner_v1/database.py

@@ -144,6 +144,7 @@ def __init__(
         self._version_retention_period = None
         self._earliest_version_time = None
         self._encryption_info = None
+        self._default_leader = None
         self.log_commit_stats = False
         self._logger = logger
         self._encryption_config = encryption_config
@@ -279,6 +280,15 @@ def encryption_info(self):
         """
         return self._encryption_info
 
+    @property
+    def default_leader(self):
+        """The read-write region which contains the database's leader replicas.
+
+        :rtype: str
+        :returns: a string representing the read-write region
+        """
+        return self._default_leader
+
     @property
     def ddl_statements(self):
         """DDL Statements used to define database schema.
@@ -414,6 +424,7 @@ def reload(self):
         self._earliest_version_time = response.earliest_version_time
         self._encryption_config = response.encryption_config
         self._encryption_info = response.encryption_info
+        self._default_leader = response.default_leader
 
     def update_ddl(self, ddl_statements, operation_id=""):
         """Update DDL for this database.

tests/system/test_system.py

@@ -68,6 +68,7 @@
     INSTANCE_ID = os.environ.get(
         "GOOGLE_CLOUD_TESTS_SPANNER_INSTANCE", "google-cloud-python-systest"
     )
+MULTI_REGION_INSTANCE_ID = "multi-region" + unique_resource_id("-")
 EXISTING_INSTANCES = []
 COUNTERS_TABLE = "counters"
 COUNTERS_COLUMNS = ("name", "value")
@@ -353,9 +354,25 @@ def setUpClass(cls):
             SPANNER_OPERATION_TIMEOUT_IN_SECONDS
         )  # raises on failure / timeout.
 
+        # Create a multi-region instance
+        multi_region_config = "nam3"
+        config_name = "{}/instanceConfigs/{}".format(
+            Config.CLIENT.project_name, multi_region_config
+        )
+        create_time = str(int(time.time()))
+        labels = {"python-spanner-systests": "true", "created": create_time}
+        cls._instance = Config.CLIENT.instance(
+            instance_id=MULTI_REGION_INSTANCE_ID,
+            configuration_name=config_name,
+            labels=labels,
+        )
+        operation = cls._instance.create()
+        operation.result(SPANNER_OPERATION_TIMEOUT_IN_SECONDS)
+
     @classmethod
     def tearDownClass(cls):
         cls._db.drop()
+        cls._instance.delete()
 
     def setUp(self):
         self.to_delete = []
@@ -443,6 +460,42 @@ def test_create_database_pitr_success(self):
             for result in results:
                 self.assertEqual(result[0], retention_period)
 
+    @unittest.skipIf(
+        USE_EMULATOR, "Default leader setting is not supported by the emulator"
+    )
+    def test_create_database_with_default_leader_success(self):
+        pool = BurstyPool(labels={"testcase": "create_database_default_leader"})
+
+        temp_db_id = "temp_db" + unique_resource_id("_")
+        default_leader = "us-east4"
+        ddl_statements = [
+            "ALTER DATABASE {}"
+            " SET OPTIONS (default_leader = '{}')".format(temp_db_id, default_leader)
+        ]
+        temp_db = self._instance.database(
+            temp_db_id, pool=pool, ddl_statements=ddl_statements
+        )
+        operation = temp_db.create()
+        self.to_delete.append(temp_db)
+
+        # We want to make sure the operation completes.
+        operation.result(30)  # raises on failure / timeout.
+
+        database_ids = [database.name for database in self._instance.list_databases()]
+        self.assertIn(temp_db.name, database_ids)
+
+        temp_db.reload()
+        self.assertEqual(temp_db.default_leader, default_leader)
+
+        with temp_db.snapshot() as snapshot:
+            results = snapshot.execute_sql(
+                "SELECT OPTION_VALUE AS default_leader "
+                "FROM INFORMATION_SCHEMA.DATABASE_OPTIONS "
+                "WHERE SCHEMA_NAME = '' AND OPTION_NAME = 'default_leader'"
+            )
+            for result in results:
+                self.assertEqual(result[0], default_leader)
+
     def test_table_not_found(self):
         temp_db_id = "temp_db" + unique_resource_id("_")
 
@@ -551,6 +604,36 @@ def test_update_database_ddl_pitr_success(self):
         self.assertEqual(temp_db.version_retention_period, retention_period)
         self.assertEqual(len(temp_db.ddl_statements), len(ddl_statements))
 
+    @unittest.skipIf(
+        USE_EMULATOR, "Default leader update is not supported by the emulator"
+    )
+    def test_update_database_ddl_default_leader_success(self):
+        pool = BurstyPool(labels={"testcase": "update_database_ddl_default_leader"})
+
+        temp_db_id = "temp_db" + unique_resource_id("_")
+        default_leader = "us-east4"
+        temp_db = self._instance.database(temp_db_id, pool=pool)
+        create_op = temp_db.create()
+        self.to_delete.append(temp_db)
+
+        # We want to make sure the operation completes.
+        create_op.result(240)  # raises on failure / timeout.
+
+        self.assertIsNone(temp_db.default_leader)
+
+        ddl_statements = DDL_STATEMENTS + [
+            "ALTER DATABASE {}"
+            " SET OPTIONS (default_leader = '{}')".format(temp_db_id, default_leader)
+        ]
+        operation = temp_db.update_ddl(ddl_statements)
+
+        # We want to make sure the operation completes.
+        operation.result(240)  # raises on failure / timeout.
+
+        temp_db.reload()
+        self.assertEqual(temp_db.default_leader, default_leader)
+        self.assertEqual(len(temp_db.ddl_statements), len(ddl_statements))
+
     def test_db_batch_insert_then_db_snapshot_read(self):
         retry = RetryInstanceState(_has_all_ddl)
         retry(self._db.reload)()

tests/unit/test_client.py

@@ -40,6 +40,7 @@ class TestClient(unittest.TestCase):
     PROCESSING_UNITS = 5000
     LABELS = {"test": "true"}
     TIMEOUT_SECONDS = 80
+    LEADER_OPTIONS = ["leader1", "leader2"]
 
     def _get_target_class(self):
         from google.cloud import spanner
@@ -457,7 +458,9 @@ def test_list_instance_configs(self):
         instance_config_pbs = ListInstanceConfigsResponse(
             instance_configs=[
                 InstanceConfigPB(
-                    name=self.CONFIGURATION_NAME, display_name=self.DISPLAY_NAME
+                    name=self.CONFIGURATION_NAME,
+                    display_name=self.DISPLAY_NAME,
+                    leader_options=self.LEADER_OPTIONS,
                 )
             ]
         )
@@ -473,6 +476,7 @@ def test_list_instance_configs(self):
         self.assertIsInstance(instance_config, InstanceConfigPB)
         self.assertEqual(instance_config.name, self.CONFIGURATION_NAME)
         self.assertEqual(instance_config.display_name, self.DISPLAY_NAME)
+        self.assertEqual(instance_config.leader_options, self.LEADER_OPTIONS)
 
         expected_metadata = (
             ("google-cloud-resource-prefix", client.project_name),

tests/unit/test_database.py

@@ -333,6 +333,13 @@ def test_encryption_info(self):
         ]
         self.assertEqual(database.encryption_info, encryption_info)
 
+    def test_default_leader(self):
+        instance = _Instance(self.INSTANCE_NAME)
+        pool = _Pool()
+        database = self._make_one(self.DATABASE_ID, instance, pool=pool)
+        default_leader = database._default_leader = "us-east4"
+        self.assertEqual(database.default_leader, default_leader)
+
     def test_spanner_api_property_w_scopeless_creds(self):
 
         client = _Client()
@@ -715,6 +722,7 @@ def test_reload_success(self):
                 kms_key_version="kms_key_version",
             )
         ]
+        default_leader = "us-east4"
         api = client.database_admin_api = self._make_database_admin_api()
         api.get_database_ddl.return_value = ddl_pb
         db_pb = Database(
@@ -725,6 +733,7 @@ def test_reload_success(self):
             earliest_version_time=_datetime_to_pb_timestamp(timestamp),
             encryption_config=encryption_config,
             encryption_info=encryption_info,
+            default_leader=default_leader,
         )
         api.get_database.return_value = db_pb
         instance = _Instance(self.INSTANCE_NAME, client=client)
@@ -740,6 +749,7 @@ def test_reload_success(self):
         self.assertEqual(database._ddl_statements, tuple(DDL_STATEMENTS))
         self.assertEqual(database._encryption_config, encryption_config)
         self.assertEqual(database._encryption_info, encryption_info)
+        self.assertEqual(database._default_leader, default_leader)
 
         api.get_database_ddl.assert_called_once_with(
             database=self.DATABASE_NAME,