diff --git a/go.mod b/go.mod index aaa5f96..48d7103 100644 --- a/go.mod +++ b/go.mod @@ -9,7 +9,7 @@ require ( github.com/Masterminds/semver v1.5.0 github.com/aserto-dev/azm v0.0.17 github.com/aserto-dev/errors v0.0.6 - github.com/aserto-dev/go-directory v0.30.5 + github.com/aserto-dev/go-directory v0.30.6 github.com/bufbuild/protovalidate-go v0.4.3 github.com/go-http-utils/headers v0.0.0-20181008091004-fed159eddc2a github.com/gonvenience/ytbx v1.4.4 diff --git a/go.sum b/go.sum index bcc08bf..782fb96 100644 --- a/go.sum +++ b/go.sum @@ -12,8 +12,8 @@ github.com/aserto-dev/azm v0.0.17 h1:jerSOk8WkNnmxt8eagT8mWSeOCswxpjQ4B1s2cZwass github.com/aserto-dev/azm v0.0.17/go.mod h1:oNNvPYysttO/otHryQDm6u71jAt/Qhs9izhrde7HlD0= github.com/aserto-dev/errors v0.0.6 h1:iH5fkJwBGFPbcdS4B8mwvNdwODlhDEXXPduZtjLh6vo= github.com/aserto-dev/errors v0.0.6/go.mod h1:kenI1gamsemaR2wS+M2un0kXIJ9exTrmeRT/fCFwlWc= -github.com/aserto-dev/go-directory v0.30.5 h1:wOF1dtiqlNbfvb7iBH9qvgQ4e2jrIfMlRJW2ly1JhSI= -github.com/aserto-dev/go-directory v0.30.5/go.mod h1:qd/+uHrKvskPSN48FLGeZ/FoiORxjRmikCKRIp3pnYY= +github.com/aserto-dev/go-directory v0.30.6 h1:JdxofQIHyqrHJxn1L/vxiYamCYFX6rGeuMUkkQkai7g= +github.com/aserto-dev/go-directory v0.30.6/go.mod h1:qd/+uHrKvskPSN48FLGeZ/FoiORxjRmikCKRIp3pnYY= github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= github.com/bufbuild/protovalidate-go v0.4.3 h1:1Xsm3qhkwioxLDEtxWgtn0Ch71xBP/sBauT/FZnn76A= github.com/bufbuild/protovalidate-go v0.4.3/go.mod h1:RcgJ+onKVv4OkAVtzkRUxkocb8stcUAMK0EoqR4fuZE= diff --git a/pkg/directory/directory.go b/pkg/directory/directory.go index 5888e74..2aa7d22 100644 --- a/pkg/directory/directory.go +++ b/pkg/directory/directory.go @@ -165,3 +165,7 @@ func (s *Directory) Reader3() dsr3.ReaderServer { func (s *Directory) Writer3() dsw3.WriterServer { return s.writer3 } + +func (s *Directory) Logger() *zerolog.Logger { + return s.logger +} diff --git a/pkg/directory/v3/exporter.go b/pkg/directory/v3/exporter.go index 99d9771..e8da071 100644 --- a/pkg/directory/v3/exporter.go +++ b/pkg/directory/v3/exporter.go @@ -1,8 +1,12 @@ package v3 import ( + "encoding/json" + "sync/atomic" + dsc3 "github.com/aserto-dev/go-directory/aserto/directory/common/v3" dse3 "github.com/aserto-dev/go-directory/aserto/directory/exporter/v3" + "github.com/aserto-dev/go-directory/pkg/pb" "github.com/aserto-dev/go-edge-ds/pkg/bdb" "github.com/rs/zerolog" @@ -26,6 +30,15 @@ func (s *Exporter) Export(req *dse3.ExportRequest, stream dse3.Exporter_ExportSe err := s.store.DB().View(func(tx *bolt.Tx) error { + // stats mode, short circuits when enabled + if req.Options&uint32(dse3.Option_OPTION_STATS) != 0 { + if err := exportStats(tx, stream, req.Options); err != nil { + logger.Error().Err(err).Msg("export_stats") + return err + } + return nil + } + if req.Options&uint32(dse3.Option_OPTION_DATA_OBJECTS) != 0 { if err := exportObjects(tx, stream); err != nil { logger.Error().Err(err).Msg("export_objects") @@ -75,3 +88,143 @@ func exportRelations(tx *bolt.Tx, stream dse3.Exporter_ExportServer) error { return nil } + +func exportStats(tx *bolt.Tx, stream dse3.Exporter_ExportServer, opts uint32) error { + stats := &Stats{ObjectTypes: ObjectTypes{}} + + // object stats. + if opts&uint32(dse3.Option_OPTION_DATA_OBJECTS) != 0 { + iter, err := bdb.NewScanIterator[dsc3.Object](stream.Context(), tx, bdb.ObjectsPath) + if err != nil { + return err + } + + for iter.Next() { + obj := iter.Value() + stats.CountObject(obj) + } + } + + // relation stats. + if opts&uint32(dse3.Option_OPTION_DATA_RELATIONS) != 0 { + iter, err := bdb.NewScanIterator[dsc3.Relation](stream.Context(), tx, bdb.RelationsObjPath) + if err != nil { + return err + } + + for iter.Next() { + rel := iter.Value() + stats.CountRelation(rel) + } + } + + buf, err := json.Marshal(stats) + if err != nil { + return err + } + + resp := pb.NewStruct() + if err := resp.UnmarshalJSON(buf); err != nil { + return err + } + + if err := stream.Send(&dse3.ExportResponse{Msg: &dse3.ExportResponse_Stats{Stats: resp}}); err != nil { + return err + } + + return nil +} + +type ObjType string +type Relation string +type SubType string +type SubRel string + +type Stats struct { + ObjectTypes ObjectTypes `json:"object_types,omitempty"` +} + +type ObjectTypes map[ObjType]struct { + ObjCount int32 `json:"_obj_count,omitempty"` + Count int32 `json:"_count,omitempty"` + Relations Relations `json:"relations,omitempty"` +} + +type Relations map[Relation]struct { + Count int32 `json:"_count,omitempty"` + SubjectTypes SubjectTypes `json:"subject_types,omitempty"` +} + +type SubjectTypes map[SubType]struct { + Count int32 `json:"_count,omitempty"` + SubjectRelations SubjectRelations `json:"subject_relations,omitempty"` +} + +type SubjectRelations map[SubRel]struct { + Count int32 `json:"_count,omitempty"` +} + +func (s *Stats) CountObject(obj *dsc3.Object) { + ot, ok := s.ObjectTypes[ObjType(obj.Type)] + if !ok { + atomic.StoreInt32(&ot.ObjCount, 0) + if ot.Relations == nil { + ot.Relations = Relations{} + } + } + + atomic.AddInt32(&ot.ObjCount, 1) + + s.ObjectTypes[ObjType(obj.Type)] = ot +} + +func (s *Stats) CountRelation(rel *dsc3.Relation) { + objType := ObjType(rel.ObjectType) + relation := Relation(rel.Relation) + subType := SubType(rel.SubjectType) + subRel := SubRel(rel.SubjectRelation) + + // object_types + ot, ok := s.ObjectTypes[objType] + if !ok { + atomic.StoreInt32(&ot.Count, 0) + } + + if ot.Relations == nil { + ot.Relations = Relations{} + } + + atomic.AddInt32(&ot.Count, 1) + s.ObjectTypes[objType] = ot + + // relations + re, ok := ot.Relations[relation] + if !ok { + atomic.StoreInt32(&re.Count, 0) + re.SubjectTypes = SubjectTypes{} + } + + atomic.AddInt32(&re.Count, 1) + ot.Relations[relation] = re + + // subject_types + st, ok := re.SubjectTypes[subType] + if !ok { + atomic.StoreInt32(&st.Count, 0) + st.SubjectRelations = SubjectRelations{} + } + + atomic.AddInt32(&st.Count, 1) + re.SubjectTypes[subType] = st + + // subject_relations + if subRel != "" { + sr, ok := st.SubjectRelations[subRel] + if !ok { + atomic.StoreInt32(&sr.Count, 0) + } + + atomic.AddInt32(&sr.Count, 1) + st.SubjectRelations[subRel] = sr + } +}