Skip to content

Commit

Permalink
Added method to join PQ and current month, with tests (#27503)
Browse files Browse the repository at this point in the history 
* Added method to join PQ and current month, with tests

* Edited test and added go doc

* Cleanup test
  • Loading branch information
@divyaac
divyaac committed Jun 18, 2024
1 parent fe1f36a commit 499d84c
Show file tree
Hide file tree
Showing 2 changed files with 325 additions and 0 deletions.
64 changes: 64 additions & 0 deletions vault/activity/query.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -319,3 +319,67 @@ func (s *PrecomputedQueryStore) DeleteQueriesBefore(ctx context.Context, retenti
}
return nil
}

func (m *MonthlyNamespaceRecord) ToNamespaceRecord() *NamespaceRecord {
return &NamespaceRecord{
NamespaceID: m.NamespaceID,
Entities: uint64(m.Counts.EntityClients),
NonEntityTokens: uint64(m.Counts.NonEntityClients),
SecretSyncs: uint64(m.Counts.SecretSyncs),
Mounts: m.Mounts,
ACMEClients: uint64(m.Counts.ACMEClients),
}
}

func (n *NamespaceRecord) CombineWithMonthlyNamespaceRecord(nsRecord *MonthlyNamespaceRecord) {
existingMounts := make(map[string]*MountRecord)
for _, mountRecord := range n.Mounts {
existingMounts[mountRecord.MountPath] = mountRecord
}

for _, mountRecord := range nsRecord.Mounts {
if existingMountRecord, ok := existingMounts[mountRecord.MountPath]; ok {
existingMountRecord.Add(mountRecord)
} else {
n.Mounts = append(n.Mounts, mountRecord)
}
}

n.SecretSyncs += uint64(nsRecord.Counts.SecretSyncs)
n.Entities += uint64(nsRecord.Counts.EntityClients)
n.NonEntityTokens += uint64(nsRecord.Counts.NonEntityClients)
n.ACMEClients += uint64(nsRecord.Counts.ACMEClients)
}

func (m *MountRecord) Add(m2 *MountRecord) {
m.Counts.ACMEClients += m2.Counts.ACMEClients
m.Counts.NonEntityClients += m2.Counts.NonEntityClients
m.Counts.EntityClients += m2.Counts.EntityClients
m.Counts.SecretSyncs += m2.Counts.SecretSyncs
}

func (q *PrecomputedQuery) CombineWithCurrentMonth(currentMonth *MonthRecord) {
// Append the current months data to the precomputed query month's data
q.Months = append(q.Months, currentMonth)

existingNamespaceMounts := make(map[string]*NamespaceRecord)
// Store the existing namespaces and mounts in the precomputed query for easy access
for _, monthlyNamespaceRecord := range q.Namespaces {
existingNamespaceMounts[monthlyNamespaceRecord.NamespaceID] = monthlyNamespaceRecord
}

// Get the counts of each mount per namespace in the current month, and increment
// its total count in the precomputed query. These total values will be visible in the
// by_namespace grouping in the final response data
for _, nsRecord := range currentMonth.Namespaces {
namespaceId := nsRecord.NamespaceID

// If the namespace already exists in the previous months, iterate through the mounts and increment the counts
if existingNsRecord, ok := existingNamespaceMounts[namespaceId]; ok {
existingNsRecord.CombineWithMonthlyNamespaceRecord(nsRecord)
} else {
// Else just add the new namespace record to the slice in the precomputed query's namespace slice
q.Namespaces = append(q.Namespaces, nsRecord.ToNamespaceRecord())
}
}
}
261 changes: 261 additions & 0 deletions vault/activity/query_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import (
"github.com/hashicorp/vault/helper/timeutil"
"github.com/hashicorp/vault/sdk/helper/logging"
"github.com/hashicorp/vault/sdk/logical"
"github.com/stretchr/testify/assert"
)

func NewTestQueryStore(t *testing.T) *PrecomputedQueryStore {
Expand Down Expand Up @@ -292,3 +293,263 @@ func TestQueryStore_TimeRanges(t *testing.T) {
})
}
}

// TestCombineWithCurrentMonth is a unit test that verifies that an
// internal join method to combine a precomputed query data structure
// with the current month data.
// This will create various repeating mounts, new mounts, and new namespaces and
// verify that these two structures are properly combined.
func TestCombineWithCurrentMonth(t *testing.T) {
// Create two months worth of records
months := []*MonthRecord{
{
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
{
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
}
pq := &PrecomputedQuery{
Months: months,
Namespaces: []*NamespaceRecord{
{
NamespaceID: "ns1",
Entities: 2,
ACMEClients: 2,
NonEntityTokens: 2,
SecretSyncs: 2,
Mounts: []*MountRecord{
{
MountPath: "m1",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
{
MountPath: "m2",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
},
},
{
NamespaceID: "ns2",
Entities: 2,
ACMEClients: 2,
NonEntityTokens: 2,
SecretSyncs: 2,
Mounts: []*MountRecord{
{
MountPath: "m1",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
{
MountPath: "m2",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
},
},
},
}

// In the new month we will add clients to three namespaces
// Namespace 1: clients for m1 (already present in PQ), and clients for m2 (already present in PQ)
// Namespace 2: clients for m1 (already present in PQ), and clients for m3 (NOT in PQ)
// Namespace 2: all clients and mounts are new
newMonthlyRecord := &MonthRecord{
Namespaces: []*MonthlyNamespaceRecord{
{
NamespaceID: "ns1",
Counts: &CountsRecord{
EntityClients: 2,
NonEntityClients: 2,
SecretSyncs: 2,
ACMEClients: 2,
},
Mounts: []*MountRecord{
{
MountPath: "m1",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
{
MountPath: "m2",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
},
},
{
NamespaceID: "ns2",
Counts: &CountsRecord{
EntityClients: 2,
NonEntityClients: 2,
SecretSyncs: 2,
ACMEClients: 2,
},
Mounts: []*MountRecord{
{
MountPath: "m1",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
{
MountPath: "m3",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
},
},
{
NamespaceID: "ns3",
Counts: &CountsRecord{
EntityClients: 2,
NonEntityClients: 2,
SecretSyncs: 2,
ACMEClients: 2,
},
Mounts: []*MountRecord{
{
MountPath: "m1",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
{
MountPath: "m2",
Counts: &CountsRecord{
EntityClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
ACMEClients: 1,
},
},
},
},
},
}

pq.CombineWithCurrentMonth(newMonthlyRecord)

// There should be 3 namespaces (one new one from this month)
assert.Equal(t, 3, len(pq.Namespaces))
// There should be 3 months (one new month from this month)
assert.Equal(t, 3, len(pq.Months))

// Verify first namespace values
assert.Equal(t, 4, int(pq.Namespaces[0].SecretSyncs))
assert.Equal(t, 4, int(pq.Namespaces[0].ACMEClients))
assert.Equal(t, 4, int(pq.Namespaces[0].Entities))
assert.Equal(t, 4, int(pq.Namespaces[0].NonEntityTokens))
assert.Equal(t, 2, len(pq.Namespaces[0].Mounts))
for i := 0; i < 2; i++ {
compareCountsRecords(t, &CountsRecord{
EntityClients: 2,
ACMEClients: 2,
NonEntityClients: 2,
SecretSyncs: 2,
}, pq.Namespaces[0].Mounts[i].Counts)
}

// Verify second namespace values
assert.Equal(t, 4, int(pq.Namespaces[1].SecretSyncs))
assert.Equal(t, 4, int(pq.Namespaces[1].ACMEClients))
assert.Equal(t, 4, int(pq.Namespaces[1].Entities))
assert.Equal(t, 4, int(pq.Namespaces[1].NonEntityTokens))
assert.Equal(t, 3, len(pq.Namespaces[1].Mounts)) // We added a new mount to this namespace (m3)
// Check that the duplicate mount got incremented properly (m1)
compareCountsRecords(t, &CountsRecord{
EntityClients: 2,
ACMEClients: 2,
NonEntityClients: 2,
SecretSyncs: 2,
}, pq.Namespaces[1].Mounts[0].Counts)
// Check the old mount counts have not changed (m2)
compareCountsRecords(t, &CountsRecord{
EntityClients: 1,
ACMEClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
}, pq.Namespaces[1].Mounts[1].Counts)
// Check the new mounts have been added (m3)
compareCountsRecords(t, &CountsRecord{
EntityClients: 1,
ACMEClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
}, pq.Namespaces[1].Mounts[2].Counts)

// Verify third namespace counts
// This is a completely new namespace
assert.Equal(t, 2, int(pq.Namespaces[2].SecretSyncs))
assert.Equal(t, 2, int(pq.Namespaces[2].ACMEClients))
assert.Equal(t, 2, int(pq.Namespaces[2].Entities))
assert.Equal(t, 2, int(pq.Namespaces[2].NonEntityTokens))
assert.Equal(t, 2, len(pq.Namespaces[2].Mounts))
compareCountsRecords(t, &CountsRecord{
EntityClients: 1,
ACMEClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
}, pq.Namespaces[2].Mounts[0].Counts)
compareCountsRecords(t, &CountsRecord{
EntityClients: 1,
ACMEClients: 1,
NonEntityClients: 1,
SecretSyncs: 1,
}, pq.Namespaces[2].Mounts[1].Counts)
}

func compareCountsRecords(t *testing.T, record *CountsRecord, toCompare *CountsRecord) {
t.Helper()
assert.Equal(t, record.NonEntityClients, toCompare.NonEntityClients)
assert.Equal(t, record.ACMEClients, toCompare.ACMEClients)
assert.Equal(t, record.SecretSyncs, toCompare.SecretSyncs)
assert.Equal(t, record.EntityClients, toCompare.EntityClients)
}

0 comments on commit 499d84c

Please sign in to comment.