group by generic way

pg-sharding · Sep 16, 2024 · f1dbfbf · f1dbfbf
1 parent 6f7dc6d
commit f1dbfbf
Show file tree

Hide file tree

Showing 6 changed files with 429 additions and 285 deletions.
diff --git a/pkg/clientinteractor/interactor.go b/pkg/clientinteractor/interactor.go
@@ -1285,24 +1285,46 @@ func (pi *PSQLInteractor) KillClient(clientID uint) error {
 //
 // Returns:
 // - error: An error if any occurred during the operation.
-func (pi *PSQLInteractor) BackendConnections(ctx context.Context, shs []shard.Shardinfo) error {
-	if err := pi.WriteHeader("backend connection id", "router", "shard key name", "hostname", "pid", "user", "dbname", "sync", "tx_served", "tx status"); err != nil {
+func (pi *PSQLInteractor) BackendConnections(ctx context.Context, shs []shard.Shardinfo, groupByClause *spqrparser.Group) error {
+	headers := []string{"backend connection id", "router", "shard key name", "hostname", "pid", "user", "dbname", "sync", "tx_served", "tx status"}
+	getters := []func(sh shard.Shardinfo) string{
+		func(sh shard.Shardinfo) string { return fmt.Sprintf("%d", sh.ID()) },
+		func(sh shard.Shardinfo) string {
+			router := "no data"
+			s, ok := sh.(shard.CoordShardinfo)
+			if ok {
+				router = s.Router()
+			}
+			return router
+		},
+		func(sh shard.Shardinfo) string { return sh.ShardKeyName() },
+		func(sh shard.Shardinfo) string { return sh.InstanceHostname() },
+		func(sh shard.Shardinfo) string { return fmt.Sprintf("%d", sh.Pid()) },
+		func(sh shard.Shardinfo) string { return sh.Usr() },
+		func(sh shard.Shardinfo) string { return sh.DB() },
+		func(sh shard.Shardinfo) string { return strconv.FormatInt(sh.Sync(), 10) },
+		func(sh shard.Shardinfo) string { return strconv.FormatInt(sh.TxServed(), 10) },
+		func(sh shard.Shardinfo) string { return sh.TxStatus().String() },
+	}
+
+	if groupByClause != nil {
+		return groupBy(headers, shs, getters, groupByClause.Col.ColName, pi)
+	}
+
+	if err := pi.WriteHeader(headers...); err != nil {
 		spqrlog.Zero.Error().Err(err).Msg("")
 		return err
 	}
 
 	for _, sh := range shs {
-		router := "no data"
-		s, ok := sh.(shard.CoordShardinfo)
-		if ok {
-			router = s.Router()
+		vals := []string{}
+		for _, getter := range getters {
+			vals = append(vals, getter(sh))
 		}
-
-		if err := pi.WriteDataRow(fmt.Sprintf("%d", sh.ID()), router, sh.ShardKeyName(), sh.InstanceHostname(), fmt.Sprintf("%d", sh.Pid()), sh.Usr(), sh.DB(), strconv.FormatInt(sh.Sync(), 10), strconv.FormatInt(sh.TxServed(), 10), sh.TxStatus().String()); err != nil {
+		if err := pi.WriteDataRow(vals...); err != nil {
 			spqrlog.Zero.Error().Err(err).Msg("")
 			return err
 		}
-
 	}
 
 	return pi.CompleteMsg(len(shs))
@@ -1390,3 +1412,29 @@ func (pi *PSQLInteractor) PreparedStatements(ctx context.Context, shs []shard.Pr
 
 	return pi.CompleteMsg(len(shs))
 }
+
+func groupBy[T any](headers []string, values []T, getters []func(s T) string, groupBy string, pi *PSQLInteractor) error {
+	ind := -1
+	for i, header := range headers {
+		if header == groupBy {
+			if err := pi.WriteHeader(groupBy, "count"); err != nil {
+				return err
+			}
+			ind = i
+			break
+		}
+	}
+
+	cnt := make(map[string]int)
+	for _, value := range values {
+		cnt[getters[ind](value)]++
+	}
+
+	for k, v := range cnt {
+		if err := pi.WriteDataRow(k, fmt.Sprintf("%d", v)); err != nil {
+			return err
+		}
+	}
+
+	return pi.CompleteMsg(len(cnt))
+}
diff --git a/pkg/meta/meta.go b/pkg/meta/meta.go
@@ -452,7 +452,11 @@ func ProcessShow(ctx context.Context, stmt *spqrparser.Show, mngr EntityMgr, ci
 			return err
 		}
 
-		return cli.BackendConnections(ctx, resp)
+		var groupBy *spqrparser.Group
+		if stmt.Group != nil {
+			groupBy = stmt.Group.(*spqrparser.Group)
+		}
+		return cli.BackendConnections(ctx, resp, groupBy)
 	case spqrparser.ShardsStr:
 		shards, err := mngr.ListShards(ctx)
 		if err != nil {

diff --git a/test/feature/features/coordinator_show.feature b/test/feature/features/coordinator_show.feature
@@ -313,6 +313,61 @@ Feature: Coordinator show clients, pools and backend_connections
         ]
         """
 
+    Scenario: show backend_connections group by
+        When I run SQL on host "coordinator"
+        """
+        SHOW backend_connections group by hostname
+        """
+        Then command return code should be "0"
+        And SQL result should match json
+        """
+        [
+            {
+                "hostname":"spqr_shard_1:6432",
+                "count": "2"
+            }
+        ]
+        """
+        And SQL result should match json
+        """
+        [
+            {
+                "hostname":"spqr_shard_2:6432",
+                "count": "2"
+            }
+        ]
+        """
+
+        When I run SQL on host "coordinator"
+        """
+        SHOW backend_connections group by user
+        """
+        Then command return code should be "0"
+        And SQL result should match json
+        """
+        [
+            {
+                "user":"regress",
+                "count": "4"
+            }
+        ]
+        """
+
+        When I run SQL on host "coordinator"
+        """
+        SHOW backend_connections group by dbname
+        """
+        Then command return code should be "0"
+        And SQL result should match json
+        """
+        [
+            {
+                "dbname":"regress",
+                "count": "4"
+            }
+        ]
+        """
+
     Scenario: show pools works
         When I run SQL on host "coordinator"
         """

diff --git a/yacc/console/ast.go b/yacc/console/ast.go
@@ -23,6 +23,13 @@ type Order struct {
 	Col        ColumnRef
 }
 
+type GroupClause interface{}
+
+type Group struct {
+	GroupClause
+	Col ColumnRef
+}
+
 type WhereClauseNode interface {
 }
 
@@ -50,6 +57,7 @@ type Show struct {
 	Cmd   string
 	Where WhereClauseNode
 	Order OrderClause
+	Group GroupClause
 }
 
 type Set struct {