Add DRZ support

src/Entity/Controller/OrganizationController.php

@@ -138,8 +138,16 @@ public function update(EntityInterface $entity): void {
   public function getEntities(): array {
     $entities = $this->decorated()->getEntities();
     foreach ($entities as $id => $entity) {
-      $this->cache[$id] = $entities;
+      $this->cache[$id] = $entity;
     }
+    return $entities;
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public function getDataResidencyEndpoint(string $organizationName): string {
+    return $this->decorated()->getDataResidencyEndpoint($organizationName);
   }
 
   /**

src/Entity/Controller/OrganizationControllerInterface.php

@@ -27,4 +27,22 @@
  */
 interface OrganizationControllerInterface extends EdgeOrganizationControllerInterface {
 
+  /**
+   * Returns the Data Residency Endpoint for the organization.
+   *
+   * This is only available for Apigee X/Hybrid organizations.
+   *
+   * @return string
+   *   Returns location based endpoint uri.
+   */
+  public function getDataResidencyEndpoint(string $organizationName): string;
+
+  /**
+   * Checks whether the organization is Edge or ApigeeX organization.
+   *
+   * @return bool
+   *   TRUE if the value of the property is "true", false otherwise.
+   */
+  public function isOrganizationApigeeX(): bool;
+
 }

src/KeyEntityFormEnhancer.php

@@ -20,6 +20,8 @@
 
 namespace Drupal\apigee_edge;
 
+use Apigee\Edge\Api\Management\Controller\OrganizationController;
+use Apigee\Edge\ClientInterface;
 use Apigee\Edge\Exception\ApiRequestException;
 use Apigee\Edge\Exception\ApigeeOnGcpOauth2AuthenticationException;
 use Apigee\Edge\Exception\OauthAuthenticationException;
@@ -345,6 +347,32 @@ public function validateForm(array &$form, FormStateInterface $form_state): void
       // Test the connection.
       $this->connector->testConnection($test_key);
       $this->messenger()->addStatus($this->t('Connection successful.'));
+      // Data Residency check.
+      $key_value_array = json_decode($key_value, TRUE);
+      if ($key_value_array['instance_type'] == EdgeKeyTypeInterface::INSTANCE_TYPE_HYBRID) {
+        try {
+          // The getProjectMapping endpoint is only available on the global endpoint.
+          $client = $this->connector->buildClient($test_key_type->getAuthenticationMethod($test_key), ClientInterface::APIGEE_ON_GCP_ENDPOINT);
+          $orgController = new OrganizationController($client);
+          $dataResidencyEndpoint = $orgController->getDataResidencyEndpoint($key_value_array['organization']);
+          if (ClientInterface::APIGEE_ON_GCP_ENDPOINT !== $dataResidencyEndpoint || ClientInterface::EDGE_ENDPOINT !== $dataResidencyEndpoint) {
+            $this->messenger()->addStatus($this->t('Data residency is enabled for this organization. Service endpoint being used is @serviceEndpoint', ['@serviceEndpoint' => $dataResidencyEndpoint]));
+            $key_value_array['drzlocation'] = $dataResidencyEndpoint;
+            $form_state->setValue('drzlocation', $dataResidencyEndpoint);
+          }
+          else {
+            unset($key_value_array['drzlocation']);
+          }
+        }
+        catch (\Exception $e) {
+          $this->messenger()->addError($this->t('Could not determine data residency information. Error: @error', ['@error' => $e->getMessage()]));
+          unset($key_value_array['drzlocation']);
+        }
+      }
+      else {
+        unset($key_value_array['drzlocation']);
+      }
+      $form_state->setValues(['key_value' => json_encode(array_filter($key_value_array))]);
 
       // Based on type of organization, cache needs to clear.
       drupal_flush_all_caches();

src/Plugin/EdgeKeyTypeBase.php

@@ -70,10 +70,13 @@ public function getAuthenticationType(KeyInterface $key): string {
    */
   public function getEndpoint(KeyInterface $key): string {
     if ($this->getInstanceType($key) === EdgeKeyTypeInterface::INSTANCE_TYPE_HYBRID) {
+      if (isset($key->getKeyValues()['drzlocation'])) {
+        return $key->getKeyValues()['drzlocation'];
+      }
       return ClientInterface::APIGEE_ON_GCP_ENDPOINT;
     }
     elseif ($this->getInstanceType($key) === EdgeKeyTypeInterface::INSTANCE_TYPE_PUBLIC) {
-      return Client::EDGE_ENDPOINT;
+      return ClientInterface::EDGE_ENDPOINT;
     }
     return $key->getKeyValues()['endpoint'];
   }

src/Plugin/KeyInput/ApigeeAuthKeyInput.php

@@ -181,6 +181,9 @@ public function buildConfigurationForm(array $form, FormStateInterface $form_sta
         'required' => $state_for_private,
       ],
     ];
+    $form['drzlocation'] = [
+      '#type' => 'hidden'
+    ];
     $form['authorization_server_type'] = [
       '#title' => $this->t('Authorization server'),
       '#type' => 'radios',
@@ -313,6 +316,9 @@ public function processSubmittedKeyValue(FormStateInterface $form_state) {
         if (!empty($input_values['gcp_hosted'])) {
           unset($input_values['gcp_hosted']);
         }
+        if (!empty($input_values['drzlocation'])) {
+          unset($input_values['drzlocation']);
+        }
         // If password field is empty we just skip it and preserve the initial
         // password if there is one already.
         if (empty($input_values['password']) && !empty($form_state->get('key_value')['current'])) {