From 4c9ba980f56b31eae2852dbd18fb6650ecab2d91 Mon Sep 17 00:00:00 2001
From: Elim Tsiagbey <elim.tsiagbey@hashicorp.com>
Date: Fri, 9 Feb 2024 14:54:10 -0500
Subject: [PATCH] add connection count info

---
 internal/cmd/commands/connect/connect.go | 13 ++++++++-----
 internal/cmd/commands/connect/funcs.go   |  3 ++-
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/internal/cmd/commands/connect/connect.go b/internal/cmd/commands/connect/connect.go
index 4234eaa6e4..c80f528923 100644
--- a/internal/cmd/commands/connect/connect.go
+++ b/internal/cmd/commands/connect/connect.go
@@ -45,7 +45,8 @@ type SessionInfo struct {
 }
 
 type ConnectionInfo struct {
-	ConnectionsLeft int32 `json:"connections_left"`
+	ConnectionsLeft  int32 `json:"connections_left"`
+	ConnectionsCount int32 `json:"connections_count"`
 }
 
 type TerminationInfo struct {
@@ -452,11 +453,12 @@ func (c *Command) Run(args []string) (retCode int) {
 				// done it manually
 				return
 			case connsLeft := <-connsLeftCh:
-				c.updateConnsLeft(connsLeft)
+				connectionsCount := clientProxy.ConnectionsCount()
+				c.updateConnsLeft(connsLeft, connectionsCount)
 
 				// If there are no available connections left and there are no connections
 				// we can exit
-				if connsLeft == 0 && clientProxy.ConnectionsCount() == 0 {
+				if connsLeft == 0 && connectionsCount == 0 {
 					return
 				}
 			}
@@ -574,9 +576,10 @@ func (c *Command) printCredentials(creds []*targets.SessionCredential) error {
 	return nil
 }
 
-func (c *Command) updateConnsLeft(connsLeft int32) {
+func (c *Command) updateConnsLeft(connsLeft int32, connsCount int32) {
 	connInfo := ConnectionInfo{
-		ConnectionsLeft: connsLeft,
+		ConnectionsLeft:  connsLeft,
+		ConnectionsCount: connsCount,
 	}
 
 	if c.flagExec == "" {
diff --git a/internal/cmd/commands/connect/funcs.go b/internal/cmd/commands/connect/funcs.go
index da115b8048..91782c48bc 100644
--- a/internal/cmd/commands/connect/funcs.go
+++ b/internal/cmd/commands/connect/funcs.go
@@ -114,7 +114,8 @@ func generateConnectionInfoTableOutput(in ConnectionInfo) string {
 	var ret []string
 
 	nonAttributeMap := map[string]any{
-		"Connections Left": in.ConnectionsLeft,
+		"Connections Left":  in.ConnectionsLeft,
+		"Connections Count": in.ConnectionsCount,
 	}
 
 	maxLength := 0