fix: add pagination to list public key mappings response (#1889)

### Proposed Changes

Missed adding pagination response to the list public key mappings
response. Also includes additional tests for limiting the result set and
also testing of filters by kas id and public key id.

### Checklist

- [x] I have added or updated unit tests
- [x] I have added or updated integration tests (if appropriate)
- [x] I have added or updated documentation

### Testing Instructions

Author: strantalis

docs/grpc/index.html

@@ -1215,6 +1215,17 @@ func (s *KasRegistrySuite) Test_List_Public_Keys_ByKasID() {
 	s.NotEqual(totalKeys, filteredTotalKeys)
 }
 
+func (s *KasRegistrySuite) Test_List_Public_Keys_WithLimit_1() {
+	r, err := s.db.PolicyClient.ListPublicKeys(s.ctx, &kasregistry.ListPublicKeysRequest{
+		Pagination: &policy.PageRequest{
+			Limit: 1,
+		},
+	})
+	s.Require().NoError(err)
+	s.NotNil(r)
+	s.Len(r.GetKeys(), 1)
+}
+
 func (s *KasRegistrySuite) Test_List_Public_Keys_WithNonExistentKasID() {
 	r, err := s.db.PolicyClient.ListPublicKeys(s.ctx, &kasregistry.ListPublicKeysRequest{KasId: nonExistentKasRegistryID})
 	s.Require().NoError(err)
@@ -1249,6 +1260,42 @@ func (s *KasRegistrySuite) Test_List_Public_Key_Mappings() {
 	}
 }
 
+func (s *KasRegistrySuite) Test_List_Public_Key_Mappings_WithLimit_1() {
+	r, err := s.db.PolicyClient.ListPublicKeyMappings(s.ctx, &kasregistry.ListPublicKeyMappingRequest{
+		Pagination: &policy.PageRequest{
+			Limit: 1,
+		},
+	})
+	s.Require().NoError(err)
+	s.NotNil(r)
+	s.Len(r.GetPublicKeyMappings(), 1)
+}
+
+func (s *KasRegistrySuite) Test_List_Public_Key_Mappings_By_KAS_ID() {
+	kasID := s.f.GetKasRegistryKey("key_access_server_1").ID
+	r, err := s.db.PolicyClient.ListPublicKeyMappings(s.ctx, &kasregistry.ListPublicKeyMappingRequest{KasId: kasID})
+	s.Require().NoError(err)
+	s.NotNil(r)
+	s.Len(r.GetPublicKeyMappings(), 1)
+
+	for _, m := range r.GetPublicKeyMappings() {
+		s.Equal(kasID, m.GetKasId())
+	}
+}
+
+func (s *KasRegistrySuite) Test_List_Public_Key_Mapping_By_PublicKey_ID() {
+	id := s.f.GetPublicKey("key_1").ID
+	r, err := s.db.PolicyClient.ListPublicKeyMappings(s.ctx, &kasregistry.ListPublicKeyMappingRequest{PublicKeyId: id})
+	s.Require().NoError(err)
+	s.NotNil(r)
+
+	for _, m := range r.GetPublicKeyMappings() {
+		s.True(slices.ContainsFunc(m.GetPublicKeys(), func(key *kasregistry.ListPublicKeyMappingResponse_PublicKey) bool {
+			return key.GetKey().GetId() == id
+		}))
+	}
+}
+
 func (s *KasRegistrySuite) Test_Deactivate_Public_Key() {
 	id := s.f.GetPublicKey("key_4").ID
 	r, err := s.db.PolicyClient.DeactivatePublicKey(s.ctx, &kasregistry.DeactivatePublicKeyRequest{Id: id})

docs/openapi/policy/kasregistry/key_access_server_registry.swagger.json

@@ -414,8 +414,9 @@ func (c PolicyDBClient) ListPublicKeyMappings(ctx context.Context, r *kasregistr
 		KasID: r.GetKasId(),
 		// KasUri:  r.GetKasUri(),
 		// KasName: r.GetKasName(),
-		Offset: offset,
-		Limit:  limit,
+		PublicKeyID: r.GetPublicKeyId(),
+		Offset:      offset,
+		Limit:       limit,
 	}
 
 	listRows, err := c.Queries.listPublicKeyMappings(ctx, params)
@@ -426,15 +427,27 @@ func (c PolicyDBClient) ListPublicKeyMappings(ctx context.Context, r *kasregistr
 	mappings := make([]*kasregistry.ListPublicKeyMappingResponse_PublicKeyMapping, len(listRows))
 	for i, mapping := range listRows {
 		pkm := new(kasregistry.ListPublicKeyMappingResponse_PublicKeyMapping)
-		err := protojson.Unmarshal(mapping, pkm)
+		err := protojson.Unmarshal(mapping.KasInfo, pkm)
 		if err != nil {
 			return nil, db.WrapIfKnownInvalidQueryErr(err)
 		}
 		mappings[i] = pkm
 	}
 
+	var total int32
+	var nextOffset int32
+	if len(listRows) > 0 {
+		total = int32(listRows[0].Total)
+		nextOffset = getNextOffset(offset, limit, total)
+	}
+
 	return &kasregistry.ListPublicKeyMappingResponse{
 		PublicKeyMappings: mappings,
+		Pagination: &policy.PageResponse{
+			CurrentOffset: params.Offset,
+			Total:         total,
+			NextOffset:    nextOffset,
+		},
 	}, nil
 }
 

protocol/go/policy/kasregistry/key_access_server_registry.pb.go

@@ -984,7 +984,12 @@ LIMIT @limit_
 OFFSET @offset_; 
 
 -- name: listPublicKeyMappings :many
-WITH base_keys AS (
+WITH counted AS (
+    SELECT COUNT(pk.id) AS total FROM public_keys AS pk
+    WHERE (NULLIF(@kas_id, '') IS NULL OR pk.key_access_server_id = @kas_id::uuid)
+    AND   ( NULLIF(@public_key_id, '') IS NULL OR pk.id = @public_key_id::uuid )
+),
+base_keys AS (
     SELECT 
         pk.id,
         pk.is_active,
@@ -1076,9 +1081,11 @@ SELECT jsonb_build_object(
             )
         )
     )
-) as kas_info
+) as kas_info,
+counted.total
 FROM base_keys bk
-GROUP BY bk.kas_id, bk.kas_name, bk.kas_uri
+CROSS JOIN counted
+GROUP BY bk.kas_id, bk.kas_name, bk.kas_uri, counted.total
 LIMIT @limit_ 
 OFFSET @offset_;
 

service/integration/kas_registry_test.go

@@ -146,6 +146,7 @@ message ListPublicKeysRequest {
   // string kas_name = 2;
   // // Optional
   // string kas_uri = 3;
+
   // Optional
   policy.PageRequest pagination = 10;
 }
@@ -164,6 +165,10 @@ message ListPublicKeyMappingRequest {
   // string kas_name = 2;
   // // Optional
   // string kas_uri = 3;
+
+  // Optional Public Key ID
+  string public_key_id = 4;
+  
   // Optional
   policy.PageRequest pagination = 10;
 }
@@ -188,6 +193,8 @@ message ListPublicKeyMappingResponse {
   }
 
   repeated PublicKeyMapping public_key_mappings = 1;
+  
+  policy.PageResponse pagination = 10;
 }
 
 message UpdatePublicKeyRequest {