authorize: performance improvements (#3723)

This commit is contained in:
Caleb Doxsey 2022-11-04 17:09:52 -06:00 committed by GitHub
parent a3cfe8fa42
commit 02df20f10a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 50 additions and 20 deletions

View file

@ -13,11 +13,13 @@ import (
opastorage "github.com/open-policy-agent/opa/storage"
"github.com/open-policy-agent/opa/storage/inmem"
"github.com/open-policy-agent/opa/types"
octrace "go.opencensus.io/trace"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/timestamppb"
"github.com/pomerium/pomerium/config"
"github.com/pomerium/pomerium/internal/log"
"github.com/pomerium/pomerium/internal/telemetry/trace"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/storage"
)
@ -105,15 +107,20 @@ func (s *Store) GetDataBrokerRecordOption() func(*rego.Rego) {
types.NewObject(nil, types.NewDynamicProperty(types.S, types.S)),
),
}, func(bctx rego.BuiltinContext, op1 *ast.Term, op2 *ast.Term) (*ast.Term, error) {
ctx, span := trace.StartSpan(bctx.Context, "rego.get_databroker_record")
defer span.End()
recordType, ok := op1.Value.(ast.String)
if !ok {
return nil, fmt.Errorf("invalid record type: %T", op1)
}
span.AddAttributes(octrace.StringAttribute("record_type", recordType.String()))
value, ok := op2.Value.(ast.String)
if !ok {
return nil, fmt.Errorf("invalid record id: %T", op2)
}
span.AddAttributes(octrace.StringAttribute("record_id", value.String()))
req := &databroker.QueryRequest{
Type: string(recordType),
@ -121,9 +128,9 @@ func (s *Store) GetDataBrokerRecordOption() func(*rego.Rego) {
}
req.SetFilterByIDOrIndex(string(value))
res, err := storage.GetQuerier(bctx.Context).Query(bctx.Context, req)
res, err := storage.GetQuerier(ctx).Query(ctx, req)
if err != nil {
log.Error(bctx.Context).Err(err).Msg("authorize/store: error retrieving record")
log.Error(ctx).Err(err).Msg("authorize/store: error retrieving record")
return ast.NullTerm(), nil
}
@ -147,7 +154,7 @@ func (s *Store) GetDataBrokerRecordOption() func(*rego.Rego) {
regoValue, err := ast.InterfaceToValue(obj)
if err != nil {
log.Error(bctx.Context).Err(err).Msg("authorize/store: error converting object to rego")
log.Error(ctx).Err(err).Msg("authorize/store: error converting object to rego")
return ast.NullTerm(), nil
}