[CLI]: Use instance/replica limits from the limits API (#2737)

Synced from monorepo@f32f60aad9793428a57d314a2052011bb144ca24

Author: SferaDev

.sync-commit

@@ -1 +1 @@
-89a52d932a3abd4ea053b95cfd51f110e739d138
+f32f60aad9793428a57d314a2052011bb144ca24

src/commands/branch/create.ts

@@ -4,6 +4,7 @@ import chalk from 'chalk';
 import { match } from 'ts-pattern';
 import type { LocalContext } from '~/context';
 import { branchConfig } from '~/lib/branch-config';
+import { getBranchLimits, instanceTypeUnavailableMessage, replicaChoicesFor } from '~/lib/branch-limits';
 import { CLI_NAME, DEFAULT_API_BASE_URL } from '~/lib/constants';
 
 import type { Types } from '@xata.io/api';
@@ -21,7 +22,7 @@ type Flags = {
   project?: string;
   name?: string;
   'parent-branch'?: string;
-  replicas?: '0' | '1' | '2' | '3' | '4';
+  replicas?: string;
   'instance-type'?: string;
   region?: string;
   'postgres-version'?: string;
@@ -30,14 +31,6 @@ type Flags = {
   json: boolean;
 };
 
-export const replicaChoices = [
-  { name: '0', message: '0' },
-  { name: '1', message: '1' },
-  { name: '2', message: '2' },
-  { name: '3', message: '3' },
-  { name: '4', message: '4' }
-];
-
 export async function instanceTypes(context: LocalContext, organizationId: string, region: string) {
   const instanceTypes = await context.api.projects.listInstanceTypes({
     pathParams: { organizationID: organizationId },
@@ -54,14 +47,17 @@ export function shouldShowInstanceTypePricing(context: LocalContext) {
 
 export function buildInstanceTypeChoices(
   instances: Awaited<ReturnType<typeof instanceTypes>>,
-  { showPricing }: { showPricing: boolean }
+  { showPricing, maxAllowedVCPUs }: { showPricing: boolean; maxAllowedVCPUs?: number }
 ) {
   return instances.map((instanceType) => {
     const parts = [`${instanceType.name} / ${instanceType.vcpus} milli-vCPU / ${instanceType.ram} GB RAM`];
     if (showPricing) {
       const instanceMonthlyCost = monthlyComputeCost(instanceType, 1);
       parts.push(chalk.gray(`$${instanceMonthlyCost.display} per mo`));
     }
+    if (maxAllowedVCPUs !== undefined && instanceType.vcpus > maxAllowedVCPUs) {
+      parts.push(chalk.yellow('not available on your current plan'));
+    }
     return {
       name: instanceType.name,
       message: parts.join(' / ')
@@ -159,13 +155,13 @@ export async function getRegion(context: LocalContext, flags: { region?: string
 
 export async function getReplicas(context: LocalContext, flags: { replicas?: string }, options: ProjectOptions) {
   const title = options?.title || 'Please select number of replicas for the branch';
+  const { maxReplicas } = await getBranchLimits(context, options.organizationId);
+  const choices = replicaChoicesFor(maxReplicas);
 
   if (flags.replicas) {
-    if (!replicaChoices.some((replica) => flags.replicas === replica.name)) {
+    if (!choices.some((replica) => flags.replicas === replica.name)) {
       context.process.stderr.write(
-        chalk.red(
-          `Invalid replica count: ${flags.replicas}. Must be one of: ${replicaChoices.map((r) => r.name).join(', ')}.`
-        )
+        chalk.red(`Invalid replica count: ${flags.replicas}. Must be one of: ${choices.map((r) => r.name).join(', ')}.`)
       );
       context.process.exit(1);
     }
@@ -174,11 +170,7 @@ export async function getReplicas(context: LocalContext, flags: { replicas?: str
   }
 
   if (!flags.replicas) {
-    const replicas = (await context.enquirer.selectPrompt(
-      context.isInteractive,
-      title,
-      replicaChoices
-    )) as Flags['replicas'];
+    const replicas = (await context.enquirer.selectPrompt(context.isInteractive, title, choices)) as Flags['replicas'];
     invariant(replicas, `Replicas should exist`);
     return replicas;
   }
@@ -192,33 +184,32 @@ export async function getInstanceType(
   options: ProjectOptions & { region: string }
 ) {
   const title = options?.title || 'Please select the type of instance for this branch';
+  const { maxAllowedVCPUs } = await getBranchLimits(context, options.organizationId);
   const instances = await instanceTypes(context, options.organizationId, options.region);
-  const instanceChoices = buildInstanceTypeChoices(instances, { showPricing: shouldShowInstanceTypePricing(context) });
+  const instanceChoices = buildInstanceTypeChoices(instances, {
+    showPricing: shouldShowInstanceTypePricing(context),
+    maxAllowedVCPUs
+  });
 
-  if (flags['instance-type']) {
-    if (!instances.some((instance) => flags['instance-type'] === instance.name)) {
-      context.process.stderr.write(
-        chalk.red(
-          `Invalid instance type: ${flags['instance-type']}. This instance type is not available for this organization.`
-        )
-      );
-      context.process.exit(1);
-    }
-    invariant(flags['instance-type'], `Instance type should exist`);
-    return flags['instance-type'];
+  const instanceType =
+    flags['instance-type'] ??
+    ((await context.enquirer.selectPrompt(context.isInteractive, title, instanceChoices)) as Flags['instance-type']);
+  invariant(instanceType, `Instance type should exist`);
+
+  const instance = instances.find((option) => option.name === instanceType);
+  if (!instance) {
+    context.process.stderr.write(
+      chalk.red(`Invalid instance type: ${instanceType}. This instance type is not available for this organization.`)
+    );
+    context.process.exit(1);
   }
 
-  if (!flags['instance-type']) {
-    const instanceType = (await context.enquirer.selectPrompt(
-      context.isInteractive,
-      title,
-      instanceChoices
-    )) as Flags['instance-type'];
-    invariant(instanceType, `Instance type should exist`);
-    return instanceType;
+  if (maxAllowedVCPUs && instance.vcpus > maxAllowedVCPUs) {
+    context.process.stderr.write(chalk.red(`${instanceTypeUnavailableMessage(instanceType)}\n`));
+    context.process.exit(1);
   }
 
-  invariant(false, `Expected input for flag --instance-type`);
+  return instanceType;
 }
 
 export async function getImage(
@@ -445,8 +436,8 @@ export const BranchCreateCommand = buildCommand({
         optional: true
       },
       replicas: {
-        kind: 'enum',
-        values: ['0', '1', '2', '3', '4'],
+        kind: 'parsed',
+        parse: String,
         brief: 'Please select number of replicas for the branch',
         optional: true
       },

src/commands/branch/set.ts

@@ -4,7 +4,8 @@ import chalk from 'chalk';
 import { match } from 'ts-pattern';
 import type { LocalContext } from '~/context';
 import { CLI_NAME } from '~/lib/constants';
-import { buildInstanceTypeChoices, instanceTypes, replicaChoices, shouldShowInstanceTypePricing } from './create';
+import { getBranchLimits, instanceTypeUnavailableMessage, replicaChoicesFor } from '~/lib/branch-limits';
+import { buildInstanceTypeChoices, instanceTypes, shouldShowInstanceTypePricing } from './create';
 import { validScaleToZeroValues, validInactivityPeriodValues, scaleToZeroChoices, timeChoices } from '~/lib/config';
 
 type Flags = {
@@ -79,8 +80,13 @@ export async function implementation(this: LocalContext, flags: Flags, fieldArg:
   const defaultScaleToZero = isRootBranch ? scaleToZeroBase : scaleToZeroChild;
   const defaultInactivityPeriod = isRootBranch ? inactivityPeriodBase : inactivityPeriodChild;
 
+  const { maxReplicas, maxAllowedVCPUs } = await getBranchLimits(this, organizationId);
+  const replicaChoices = replicaChoicesFor(maxReplicas);
   const instances = await instanceTypes(this, organizationId, branchRegion);
-  const instanceChoices = buildInstanceTypeChoices(instances, { showPricing: shouldShowInstanceTypePricing(this) });
+  const instanceChoices = buildInstanceTypeChoices(instances, {
+    showPricing: shouldShowInstanceTypePricing(this),
+    maxAllowedVCPUs
+  });
 
   let upgradeableImageChoices: { name: string; message: string }[] = [];
   if (field === 'postgres-version') {
@@ -172,7 +178,7 @@ export async function implementation(this: LocalContext, flags: Flags, fieldArg:
       }
     })
     .with('replicas', () => {
-      const validReplicas = ['0', '1', '2', '3', '4'];
+      const validReplicas = replicaChoices.map((choice) => choice.name);
       if (!validReplicas.includes(value)) {
         this.process.stderr.write(
           chalk.red(`Invalid replicas value: ${value}. Valid values are: ${validReplicas.join(', ')}`)
@@ -181,13 +187,18 @@ export async function implementation(this: LocalContext, flags: Flags, fieldArg:
       }
     })
     .with('instance-type', () => {
-      const validInstanceTypes = instances.map((t) => t.name);
-      if (!validInstanceTypes.includes(value)) {
+      const instance = instances.find((t) => t.name === value);
+      if (!instance) {
+        const validInstanceTypes = instances.map((t) => t.name);
         this.process.stderr.write(
           chalk.red(`Invalid instance type: ${value}. Valid values are: ${validInstanceTypes.join(', ')}`)
         );
         this.process.exit(1);
       }
+      if (maxAllowedVCPUs && instance.vcpus > maxAllowedVCPUs) {
+        this.process.stderr.write(chalk.red(`${instanceTypeUnavailableMessage(value)}\n`));
+        this.process.exit(1);
+      }
     })
     .with('hibernate', () => {
       const validHibernateValues = ['true', 'false'];

src/commands/onboard.ts

@@ -13,7 +13,7 @@ type Flags = {
   'project-name'?: string;
   'branch-name'?: string;
 
-  replicas?: '0' | '1' | '2' | '3' | '4';
+  replicas?: string;
   'instance-type'?: string;
   region?: string;
   'scale-to-zero-base'?: 'true' | 'false';
@@ -229,8 +229,8 @@ export const OnboardCommand = buildCommand({
         optional: true
       },
       replicas: {
-        kind: 'enum',
-        values: ['0', '1', '2', '3', '4'],
+        kind: 'parsed',
+        parse: String,
         brief: 'Please select number of replicas for the branch',
         optional: true
       },

src/commands/project/create.ts

@@ -7,7 +7,7 @@ type Flags = {
   organization?: string;
   name: string;
   'branch-name'?: string;
-  replicas?: '0' | '1' | '2' | '3' | '4';
+  replicas?: string;
   'instance-type'?: string;
   region?: string;
   'scale-to-zero-base'?: 'true' | 'false';
@@ -144,8 +144,8 @@ export const ProjectCreateCommand = buildCommand({
         optional: true
       },
       replicas: {
-        kind: 'enum',
-        values: ['0', '1', '2', '3', '4'],
+        kind: 'parsed',
+        parse: String,
         brief: 'Please select number of replicas for the branch',
         optional: true
       },

src/lib/branch-limits.ts

@@ -0,0 +1,33 @@
+import type { LocalContext } from '~/context';
+
+const DEFAULT_MAX_INSTANCES_PER_BRANCH = 5;
+
+export type BranchLimits = { maxReplicas: number; maxAllowedVCPUs: number | undefined };
+
+// Effective instance/replica limits for the org. Falls back to today's defaults
+// (4 replicas, no instance-type cap) if the limits endpoint is unavailable, so
+// commands never fail on limits alone.
+export async function getBranchLimits(context: LocalContext, organizationId: string): Promise<BranchLimits> {
+  try {
+    const limits = await context.api.projects.getOrganizationLimits({
+      pathParams: { organizationID: organizationId }
+    });
+    return {
+      maxReplicas: (limits.maxInstancesPerBranch ?? DEFAULT_MAX_INSTANCES_PER_BRANCH) - 1,
+      maxAllowedVCPUs: limits.maxAllowedInstanceType
+    };
+  } catch {
+    return { maxReplicas: DEFAULT_MAX_INSTANCES_PER_BRANCH - 1, maxAllowedVCPUs: undefined };
+  }
+}
+
+export function replicaChoicesFor(maxReplicas: number) {
+  return Array.from({ length: Math.max(0, maxReplicas) + 1 }, (_, count) => ({
+    name: String(count),
+    message: String(count)
+  }));
+}
+
+export function instanceTypeUnavailableMessage(instanceType: string): string {
+  return `Instance type ${instanceType} is not available on your current plan; please add a payment method in your billing settings or contact support to enable larger instances.`;
+}