Connect using private IPs for Cloud SQL.

PiperOrigin-RevId: 715950195

Author: bvliu

perfkitbenchmarker/providers/gcp/gcp_alloy_db.py

@@ -108,6 +108,11 @@ def GetResourceMetadata(self) -> Dict[str, Any]:
     })
     return metadata
 
+  def _CreateDependencies(self):
+    util.SetupPrivateServicesAccess(
+        self.client_vm.network.network_resource.name, self.project
+    )
+
   def _Create(self) -> None:
     """Creates the Cloud SQL instance and authorizes traffic from anywhere.
 

perfkitbenchmarker/providers/gcp/gcp_relational_db.py

@@ -165,13 +165,15 @@ class GCPOmniPostgresIAASRelationalDb(
     omni_postgres_iaas_relational_db.OmniPostgresIAASRelationalDb
 ):
   """A GCP Omni Postgres IAAS database resource."""
+
   CLOUD = provider_info.GCP
 
 
 class GCPTimescaleDbPostgresIAASRelationalDb(
     timescaledb_iaas_relational_db.TimescaleDbIAASRelationalDb
 ):
   """A TimescaleDB Postgres IAAS database resource."""
+
   CLOUD = provider_info.GCP
 
 
@@ -191,24 +193,15 @@ def __init__(self, relational_db_spec):
     super().__init__(relational_db_spec)
     self.project = FLAGS.project or util.GetDefaultProject()
 
-  def _GetAuthorizedNetworks(self, vms):
-    """Get CIDR connections for list of VM specs that need to access the db."""
-    for vm in vms:
-      if not vm.HasIpAddress:
-        raise RuntimeError(
-            'Client vm needs to be initialized before database can '
-            'discover authorized network.'
-        )
-    # create the CIDR of the client VM that is configured to access
-    # the database
-    return ','.join('{}/32'.format(vm.ip_address) for vm in vms)
+  def _CreateDependencies(self):
+    util.SetupPrivateServicesAccess(
+        self.client_vm.network.network_resource.name, self.project
+    )
 
   def _CreateGcloudSqlInstance(self):
     storage_size = self.spec.db_disk_spec.disk_size
     instance_zone = self.spec.db_spec.zone
 
-    authorized_network = self._GetAuthorizedNetworks([self.client_vm])
-
     database_version_string = self._GetEngineVersionString(
         self.spec.engine, self.spec.engine_version
     )
@@ -222,8 +215,9 @@ def _CreateGcloudSqlInstance(self):
         '--quiet',
         '--format=json',
         '--activation-policy=ALWAYS',
-        '--assign-ip',
-        '--authorized-networks=%s' % authorized_network,
+        '--no-assign-ip',
+        '--network=%s' % self.client_vm.network.network_resource.name,
+        '--allocated-ip-range-name=google-service-range',
         '--zone=%s' % instance_zone,
         '--database-version=%s' % database_version_string,
         '--storage-size=%d' % storage_size,
@@ -584,8 +578,9 @@ def _GetEngineVersionString(engine, version):
     if engine not in GCP_DATABASE_VERSION_MAPPING:
       valid_databases = ', '.join(GCP_DATABASE_VERSION_MAPPING.keys())
       raise NotImplementedError(
-          'Database {} is not supported,supported '
-          'databases include {}'.format(engine, valid_databases)
+          'Database {} is not supported,supported databases include {}'.format(
+              engine, valid_databases
+          )
       )
 
     version_mapping = GCP_DATABASE_VERSION_MAPPING[engine]

perfkitbenchmarker/providers/gcp/util.py

@@ -543,9 +543,7 @@ def FormatTags(tags_dict: dict[str, str]):
   Returns:
     A string contains formatted tags
   """
-  return ','.join(
-      '{}={}'.format(k, v) for k, v in sorted(tags_dict.items())
-  )
+  return ','.join('{}={}'.format(k, v) for k, v in sorted(tags_dict.items()))
 
 
 def SplitTags(tags: str):
@@ -604,3 +602,52 @@ def GetAccessToken(application_default: bool = True) -> str:
   cmd.append('print-access-token')
   stdout, _, _ = vm_util.IssueCommand(cmd)
   return stdout.strip()
+
+
+def SetupPrivateServicesAccess(network: str, project: str) -> None:
+  """Setup private services access for the network.
+
+  Private services access is used by several GCP services and enables VMs to
+  connect to service producers using internal IP addressses. See
+  https://cloud.google.com/vpc/docs/configure-private-services-access for more
+  info.
+
+  Args:
+    network: The network (VPC) name to setup private services access for.
+    project: The project to setup private services access for.
+  """
+  cmd = GcloudCommand(
+      None,
+      'compute',
+      'addresses',
+      'create',
+      'google-service-range',
+  )
+  cmd.flags['global'] = True
+  cmd.flags['prefix-length'] = 16
+  cmd.flags['purpose'] = 'VPC_PEERING'
+  cmd.flags['network'] = network
+  cmd.flags['project'] = project
+  cmd.Issue(raise_on_failure=False)
+
+  cmd = GcloudCommand(None, 'services', 'vpc-peerings', 'connect')
+  cmd.flags['service'] = 'servicenetworking.googleapis.com'
+  cmd.flags['network'] = network
+  cmd.flags['ranges'] = 'google-service-range'
+  cmd.flags['project'] = project
+  _, stderr, _ = cmd.Issue(raise_on_failure=False)
+
+  # There are create errors when creating the connection for a second time.
+  # The workaround is to update the connection:
+  if 'Please use UpdateConnection' in stderr:
+    cmd = GcloudCommand(
+        None,
+        'services',
+        'vpc-peerings',
+        'update',
+    )
+    cmd.flags['network'] = network
+    cmd.flags['project'] = project
+    cmd.flags['ranges'] = 'google-service-range'
+    cmd.flags['force'] = True
+    cmd.Issue(raise_on_failure=False)