1 files changed, 51 insertions, 15 deletions
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/LockedTenant.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/LockedTenant.java
index 7d19acfce80..bfba17bef22 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/LockedTenant.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/LockedTenant.java
@@ -1,4 +1,4 @@
-// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller;
 
 import com.google.common.collect.BiMap;
@@ -10,6 +10,7 @@ import com.yahoo.transaction.Mutex;
 import com.yahoo.vespa.athenz.api.AthenzDomain;
 import com.yahoo.vespa.hosted.controller.api.identifiers.Property;
 import com.yahoo.vespa.hosted.controller.api.identifiers.PropertyId;
+import com.yahoo.vespa.hosted.controller.api.integration.billing.PlanId;
 import com.yahoo.vespa.hosted.controller.api.integration.organization.Contact;
 import com.yahoo.vespa.hosted.controller.api.integration.secrets.TenantSecretStore;
 import com.yahoo.vespa.hosted.controller.api.role.SimplePrincipal;
@@ -151,12 +152,14 @@ public abstract class LockedTenant {
         private final ArchiveAccess archiveAccess;
         private final Optional<Instant> invalidateUserSessionsBefore;
         private final Optional<BillingReference> billingReference;
+        private final PlanId planId;
 
         private Cloud(TenantName name, Instant createdAt, LastLoginInfo lastLoginInfo, Optional<SimplePrincipal> creator,
                       BiMap<PublicKey, SimplePrincipal> developerKeys, TenantInfo info,
                       List<TenantSecretStore> tenantSecretStores, ArchiveAccess archiveAccess,
                       Optional<Instant> invalidateUserSessionsBefore, Instant tenantRolesLastMaintained,
-                      List<CloudAccountInfo> cloudAccounts, Optional<BillingReference> billingReference) {
+                      List<CloudAccountInfo> cloudAccounts, Optional<BillingReference> billingReference,
+                      PlanId planId) {
             super(name, createdAt, lastLoginInfo, tenantRolesLastMaintained, cloudAccounts);
             this.developerKeys = ImmutableBiMap.copyOf(developerKeys);
             this.creator = creator;
@@ -165,15 +168,20 @@ public abstract class LockedTenant {
             this.archiveAccess = archiveAccess;
             this.invalidateUserSessionsBefore = invalidateUserSessionsBefore;
             this.billingReference = billingReference;
+            this.planId = planId;
         }
 
         private Cloud(CloudTenant tenant) {
-            this(tenant.name(), tenant.createdAt(), tenant.lastLoginInfo(), tenant.creator(), tenant.developerKeys(), tenant.info(), tenant.tenantSecretStores(), tenant.archiveAccess(), tenant.invalidateUserSessionsBefore(), tenant.tenantRolesLastMaintained(), tenant.cloudAccounts(), tenant.billingReference());
+            this(tenant.name(), tenant.createdAt(), tenant.lastLoginInfo(), tenant.creator(), tenant.developerKeys(),
+                 tenant.info(), tenant.tenantSecretStores(), tenant.archiveAccess(), tenant.invalidateUserSessionsBefore(),
+                 tenant.tenantRolesLastMaintained(), tenant.cloudAccounts(), tenant.billingReference(), tenant.planId());
         }
 
         @Override
         public CloudTenant get() {
-            return new CloudTenant(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new CloudTenant(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores,
+                                   archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained,
+                                   cloudAccounts, billingReference, planId);
         }
 
         public Cloud withDeveloperKey(PublicKey key, Principal principal) {
@@ -184,56 +192,84 @@ public abstract class LockedTenant {
             if (keys.inverse().containsKey(simplePrincipal))
                 throw new IllegalArgumentException(principal + " is already associated with key " + KeyUtils.toPem(keys.inverse().get(simplePrincipal)));
             keys.put(key, simplePrincipal);
-            return new Cloud(name, createdAt, lastLoginInfo, creator, keys, info, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, keys, info, tenantSecretStores, archiveAccess,
+                             invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
 
         public Cloud withoutDeveloperKey(PublicKey key) {
             BiMap<PublicKey, SimplePrincipal> keys = HashBiMap.create(developerKeys);
             keys.remove(key);
-            return new Cloud(name, createdAt, lastLoginInfo, creator, keys, info, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, keys, info, tenantSecretStores, archiveAccess,
+                             invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference,
+                             planId);
         }
 
         public Cloud withInfo(TenantInfo newInfo) {
-            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, newInfo, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, newInfo, tenantSecretStores,
+                             archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
 
         @Override
         public LockedTenant with(LastLoginInfo lastLoginInfo) {
-            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores,
+                             archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
 
         public Cloud withSecretStore(TenantSecretStore tenantSecretStore) {
             ArrayList<TenantSecretStore> secretStores = new ArrayList<>(tenantSecretStores);
             secretStores.add(tenantSecretStore);
-            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, secretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, secretStores, archiveAccess,
+                             invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
 
         public Cloud withoutSecretStore(TenantSecretStore tenantSecretStore) {
             ArrayList<TenantSecretStore> secretStores = new ArrayList<>(tenantSecretStores);
             secretStores.remove(tenantSecretStore);
-            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, secretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, secretStores, archiveAccess,
+                             invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
 
         public Cloud withArchiveAccess(ArchiveAccess archiveAccess) {
-            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore,tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess,
+                             invalidateUserSessionsBefore,tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
 
         public Cloud withInvalidateUserSessionsBefore(Instant invalidateUserSessionsBefore) {
-            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess, Optional.of(invalidateUserSessionsBefore), tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess,
+                             Optional.of(invalidateUserSessionsBefore), tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
 
         @Override
         public LockedTenant with(Instant tenantRolesLastMaintained) {
-            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess,
+                             invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
 
         @Override
         public LockedTenant withCloudAccounts(List<CloudAccountInfo> cloudAccounts) {
-            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, billingReference);
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess,
+                             invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
 
         public Cloud with(BillingReference billingReference) {
-            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts, Optional.of(billingReference));
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess,
+                             invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             Optional.of(billingReference), planId);
+        }
+
+        public Cloud withPlanId(PlanId planId) {
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess,
+                             invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, planId);
         }
     }