Add support for persisting plan in ZooKeeper for a tenant

7 files changed, 103 insertions, 18 deletions

diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/tenant/CloudTenant.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/tenant/CloudTenant.java
index 173d3e1950e..ae6c15852b4 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/tenant/CloudTenant.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/tenant/CloudTenant.java
@@ -28,13 +28,15 @@ public class CloudTenant extends Tenant {
     private final ArchiveAccess archiveAccess;
     private final Optional<Instant> invalidateUserSessionsBefore;
     private final Optional<BillingReference> billingReference;
+    private final Optional<String> planId;
 
     /** Public for the serialization layer — do not use! */
     public CloudTenant(TenantName name, Instant createdAt, LastLoginInfo lastLoginInfo, Optional<SimplePrincipal> creator,
                        BiMap<PublicKey, SimplePrincipal> developerKeys, TenantInfo info,
                        List<TenantSecretStore> tenantSecretStores, ArchiveAccess archiveAccess,
                        Optional<Instant> invalidateUserSessionsBefore, Instant tenantRoleLastMaintained,
-                       List<CloudAccountInfo> cloudAccounts, Optional<BillingReference> billingReference) {
+                       List<CloudAccountInfo> cloudAccounts, Optional<BillingReference> billingReference,
+                       Optional<String> planId) {
         super(name, createdAt, lastLoginInfo, Optional.empty(), tenantRoleLastMaintained, cloudAccounts);
         this.creator = creator;
         this.developerKeys = developerKeys;
@@ -43,6 +45,7 @@ public class CloudTenant extends Tenant {
         this.archiveAccess = Objects.requireNonNull(archiveAccess);
         this.invalidateUserSessionsBefore = invalidateUserSessionsBefore;
         this.billingReference = Objects.requireNonNull(billingReference);
+        this.planId = Objects.requireNonNull(planId);
     }
 
     /** Creates a tenant with the given name, provided it passes validation. */
@@ -52,7 +55,7 @@ public class CloudTenant extends Tenant {
                                LastLoginInfo.EMPTY,
                                Optional.ofNullable(creator).map(SimplePrincipal::of),
                                ImmutableBiMap.of(), TenantInfo.empty(), List.of(), new ArchiveAccess(), Optional.empty(),
-                               Instant.EPOCH, List.of(), Optional.empty());
+                               Instant.EPOCH, List.of(), Optional.empty(), Optional.empty());
     }
 
     /** The user that created the tenant */
@@ -92,6 +95,8 @@ public class CloudTenant extends Tenant {
         return billingReference;
     }
 
+    public Optional<String> planId() { return planId; }
+
     @Override
     public Type type() {
         return Type.cloud;
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..d9854543b72 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
@@ -151,12 +151,14 @@ public abstract class LockedTenant {
         private final ArchiveAccess archiveAccess;
         private final Optional<Instant> invalidateUserSessionsBefore;
         private final Optional<BillingReference> billingReference;
+        private final Optional<String> 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,
+                      Optional<String> planId) {
             super(name, createdAt, lastLoginInfo, tenantRolesLastMaintained, cloudAccounts);
             this.developerKeys = ImmutableBiMap.copyOf(developerKeys);
             this.creator = creator;
@@ -165,15 +167,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 +191,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(String planId) {
+            return new Cloud(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess,
+                             invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccounts,
+                             billingReference, Optional.ofNullable(planId));
         }
     }
 
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/persistence/TenantSerializer.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/persistence/TenantSerializer.java
index 760fb9b0366..4021d5161c4 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/persistence/TenantSerializer.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/persistence/TenantSerializer.java
@@ -88,6 +88,7 @@ public class TenantSerializer {
     private static final String invalidateUserSessionsBeforeField = "invalidateUserSessionsBefore";
     private static final String tenantRolesLastMaintainedField = "tenantRolesLastMaintained";
     private static final String billingReferenceField = "billingReference";
+    private static final String planIdField = "planId";
     private static final String cloudAccountsField = "cloudAccounts";
     private static final String accountField = "account";
     private static final String templateVersionField = "templateVersion";
@@ -137,6 +138,7 @@ public class TenantSerializer {
         toSlime(tenant.archiveAccess(), root);
         tenant.billingReference().ifPresent(b -> toSlime(b, root));
         tenant.invalidateUserSessionsBefore().ifPresent(instant -> root.setLong(invalidateUserSessionsBeforeField, instant.toEpochMilli()));
+        tenant.planId().ifPresent(id -> root.setString(planIdField, id));
     }
 
     private void toSlime(ArchiveAccess archiveAccess, Cursor root) {
@@ -215,7 +217,10 @@ public class TenantSerializer {
         Instant tenantRolesLastMaintained = SlimeUtils.instant(tenantObject.field(tenantRolesLastMaintainedField));
         List<CloudAccountInfo> cloudAccountInfos = cloudAccountsFromSlime(tenantObject.field(cloudAccountsField));
         Optional<BillingReference> billingReference = billingReferenceFrom(tenantObject.field(billingReferenceField));
-        return new CloudTenant(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores, archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained, cloudAccountInfos, billingReference);
+        Optional<String> planId = planId(tenantObject.field(planIdField));
+        return new CloudTenant(name, createdAt, lastLoginInfo, creator, developerKeys, info, tenantSecretStores,
+                               archiveAccess, invalidateUserSessionsBefore, tenantRolesLastMaintained,
+                               cloudAccountInfos, billingReference, planId);
     }
 
     private DeletedTenant deletedTenantFrom(Inspector tenantObject) {
@@ -250,6 +255,7 @@ public class TenantSerializer {
                 .withAWSRole(awsArchiveAccessRole)
                 .withGCPMember(gcpArchiveAccessMember);
     }
+
     TenantInfo tenantInfoFromSlime(Inspector infoObject) {
         if (!infoObject.valid()) return TenantInfo.empty();
 
@@ -375,6 +381,12 @@ public class TenantSerializer {
                 SlimeUtils.instant(object.field("updated"))));
     }
 
+    private Optional<String> planId(Inspector object) {
+        if (! object.valid()) return Optional.empty();
+
+        return Optional.of(object.asString());
+    }
+
     private TenantContacts tenantContactsFrom(Inspector object) {
         List<TenantContacts.Contact> contacts = SlimeUtils.entriesStream(object)
                 .map(this::readContact)
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/notification/NotificationsDbTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/notification/NotificationsDbTest.java
index 228a61cebc6..d4c4083d1b2 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/notification/NotificationsDbTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/notification/NotificationsDbTest.java
@@ -27,6 +27,7 @@ import com.yahoo.vespa.hosted.controller.tenant.Email;
 import com.yahoo.vespa.hosted.controller.tenant.LastLoginInfo;
 import com.yahoo.vespa.hosted.controller.tenant.TenantContacts;
 import com.yahoo.vespa.hosted.controller.tenant.TenantInfo;
+import org.apache.zookeeper.Op;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -69,6 +70,7 @@ public class NotificationsDbTest {
             Optional.empty(),
             Instant.EPOCH,
             List.of(),
+            Optional.empty(),
             Optional.empty());
     private static final List<Notification> notifications = List.of(
             notification(1001, Type.deployment, Level.error, NotificationSource.from(tenant), "tenant msg"),
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/notification/NotifierTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/notification/NotifierTest.java
index 15524e2748c..6aa32d7d283 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/notification/NotifierTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/notification/NotifierTest.java
@@ -47,6 +47,7 @@ public class NotifierTest {
             Optional.empty(),
             Instant.EPOCH,
             List.of(),
+            Optional.empty(),
             Optional.empty());
 
 
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/persistence/TenantSerializerTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/persistence/TenantSerializerTest.java
index 4369675ba3e..d5f2eb0162e 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/persistence/TenantSerializerTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/persistence/TenantSerializerTest.java
@@ -114,12 +114,14 @@ public class TenantSerializerTest {
                 Optional.empty(),
                 Instant.EPOCH,
                 List.of(),
+                Optional.empty(),
                 Optional.empty());
         CloudTenant serialized = (CloudTenant) serializer.tenantFrom(serializer.toSlime(tenant));
         assertEquals(tenant.name(), serialized.name());
         assertEquals(tenant.creator(), serialized.creator());
         assertEquals(tenant.developerKeys(), serialized.developerKeys());
         assertEquals(tenant.createdAt(), serialized.createdAt());
+        assertTrue(serialized.planId().isEmpty());
     }
 
     @Test
@@ -139,6 +141,7 @@ public class TenantSerializerTest {
                 Optional.of(Instant.ofEpochMilli(1234567)),
                 Instant.EPOCH,
                 List.of(),
+                Optional.empty(),
                 Optional.empty());
         CloudTenant serialized = (CloudTenant) serializer.tenantFrom(serializer.toSlime(tenant));
         assertEquals(tenant.info(), serialized.info());
@@ -193,6 +196,7 @@ public class TenantSerializerTest {
                 Instant.EPOCH,
                 List.of(new CloudAccountInfo(CloudAccount.from("aws:123456789012"), Version.fromString("1.2.3")),
                         new CloudAccountInfo(CloudAccount.from("gcp:my-project"), Version.fromString("3.2.1"))),
+                Optional.empty(),
                 Optional.empty());
         CloudTenant serialized = (CloudTenant) serializer.tenantFrom(serializer.toSlime(tenant));
         assertEquals(serialized.archiveAccess().awsRole().get(), "arn:aws:iam::123456789012:role/my-role");
@@ -253,6 +257,30 @@ public class TenantSerializerTest {
     }
 
     @Test
+    void cloud_tenant_with_plan_id() {
+        CloudTenant tenant = new CloudTenant(TenantName.from("elderly-lady"),
+                                             Instant.ofEpochMilli(1234L),
+                                             lastLoginInfo(123L, 456L, null),
+                                             Optional.of(new SimplePrincipal("foobar-user")),
+                                             ImmutableBiMap.of(publicKey, new SimplePrincipal("joe"),
+                                                               otherPublicKey, new SimplePrincipal("jane")),
+                                             TenantInfo.empty(),
+                                             List.of(),
+                                             new ArchiveAccess(),
+                                             Optional.empty(),
+                                             Instant.EPOCH,
+                                             List.of(),
+                                             Optional.empty(),
+                                             Optional.of("pay-as-you-go"));
+        CloudTenant serialized = (CloudTenant) serializer.tenantFrom(serializer.toSlime(tenant));
+        assertEquals(tenant.name(), serialized.name());
+        assertEquals(tenant.creator(), serialized.creator());
+        assertEquals(tenant.developerKeys(), serialized.developerKeys());
+        assertEquals(tenant.createdAt(), serialized.createdAt());
+        assertEquals(tenant.planId(), serialized.planId());
+    }
+
+    @Test
     void deleted_tenant() {
         DeletedTenant tenant = new DeletedTenant(
                 TenantName.from("tenant1"), Instant.ofEpochMilli(1234L), Instant.ofEpochMilli(2345L));
@@ -291,7 +319,8 @@ public class TenantSerializerTest {
                 Optional.empty(),
                 Instant.EPOCH,
                 List.of(),
-                Optional.of(reference));
+                Optional.of(reference),
+                Optional.empty());
         var slime = serializer.toSlime(tenant);
         var deserialized = serializer.tenantFrom(slime);
         assertEquals(tenant, deserialized);
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/filter/SignatureFilterTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/filter/SignatureFilterTest.java
index 001e02e1b16..fa6f323749d 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/filter/SignatureFilterTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/filter/SignatureFilterTest.java
@@ -120,6 +120,7 @@ public class SignatureFilterTest {
                 Optional.empty(),
                 Instant.EPOCH,
                 List.of(),
+                Optional.empty(),
                 Optional.empty()));
         verifySecurityContext(requestOf(signer.signed(request.copy(), Method.POST, () -> new ByteArrayInputStream(hiBytes)), hiBytes),
                 new SecurityContext(new SimplePrincipal("user"),