Skip to content

Write usage logs when diverting to spanner #2075

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
n-h-diaz merged 4 commits into from
Jul 30, 2026
Merged

Write usage logs when diverting to spanner #2075

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
145363f
Write usage logs when diverting to spanner
n-h-diaz Jul 30, 2026
00869d9
Merge branch 'master' into usage-logs-dispatcher
n-h-diaz Jul 30, 2026
994d955
fix
n-h-diaz Jul 30, 2026
1ff226f
Merge branch 'usage-logs-dispatcher' of https://github.com/n-h-diaz/m…
n-h-diaz Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 53 additions & 44 deletions internal/server/handler_v2.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import (
"github.com/datacommonsorg/mixer/internal/server/translator"
v2observation "github.com/datacommonsorg/mixer/internal/server/v2/observation"
"github.com/datacommonsorg/mixer/internal/server/v2/resolve"
"github.com/datacommonsorg/mixer/internal/server/v2/shared"
"github.com/datacommonsorg/mixer/internal/util"
"github.com/google/uuid"
"golang.org/x/sync/errgroup"
Expand Down Expand Up @@ -369,53 +370,61 @@ func (s *Server) handleV2Event(
func (s *Server) V2Observation(
ctx context.Context, in *pbv2.ObservationRequest,
) (*pbv2.ObservationResponse, error) {
if s.shouldDivertV2(ctx) {
return s.dispatcher.Observation(ctx, in)
}

v2StartTime := time.Now()

var v2Resp *pbv2.ObservationResponse
var queryType shared.QueryType
var err error
surface, toRemote := util.GetMetadata(ctx)

initialResp, queryType, err := v2observation.ObservationInternal(
ctx,
s.store,
s.cachedata.Load(),
s.metadata,
s.httpClient,
in,
surface)
if err != nil {
return nil, err
}
calculatedResps, err := v2observation.MaybeCalculateHoles(
ctx,
s.store,
s.cachedata.Load(),
s.metadata,
s.httpClient,
in,
initialResp,
surface,
)
if err != nil {
return nil, err
if s.shouldDivertV2(ctx) {
v2Resp, err = s.dispatcher.Observation(ctx, in)
if err != nil {
return nil, err
}
queryType = v2observation.GetQueryType(in)
} else {
v2StartTime := time.Now()

initialResp, qt, err := v2observation.ObservationInternal(
ctx,
s.store,
s.cachedata.Load(),
s.metadata,
s.httpClient,
in,
surface)
if err != nil {
return nil, err
}
queryType = qt
calculatedResps, err := v2observation.MaybeCalculateHoles(
ctx,
s.store,
s.cachedata.Load(),
s.metadata,
s.httpClient,
in,
initialResp,
surface,
)
if err != nil {
return nil, err
}
// initialResp is preferred over any calculated response.
combinedResp := append([]*pbv2.ObservationResponse{initialResp}, calculatedResps...)
v2Resp = merger.MergeMultiObservation(combinedResp)
v2Latency := time.Since(v2StartTime)

s.maybeMirrorV3(
ctx,
in,
v2Resp,
v2Latency,
func(ctx context.Context, req proto.Message) (proto.Message, error) {
return s.V3Observation(ctx, req.(*pbv2.ObservationRequest))
},
GetV2ObservationCmpOpts(),
)
}
// initialResp is preferred over any calculated response.
combinedResp := append([]*pbv2.ObservationResponse{initialResp}, calculatedResps...)
v2Resp := merger.MergeMultiObservation(combinedResp)
v2Latency := time.Since(v2StartTime)

s.maybeMirrorV3(
ctx,
in,
v2Resp,
v2Latency,
func(ctx context.Context, req proto.Message) (proto.Message, error) {
return s.V3Observation(ctx, req.(*pbv2.ObservationRequest))
},
GetV2ObservationCmpOpts(),
)

// Create a new ID to return as a header on the response.
// This is used for usage logging and in the website to log cached usage.
Expand Down
102 changes: 71 additions & 31 deletions internal/server/handler_v2_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ import (
"github.com/datacommonsorg/mixer/internal/featureflags"
pbv2 "github.com/datacommonsorg/mixer/internal/proto/v2"
"github.com/datacommonsorg/mixer/internal/server/cache"
"github.com/datacommonsorg/mixer/internal/server/datasource"
"github.com/datacommonsorg/mixer/internal/server/datasources"
"github.com/datacommonsorg/mixer/internal/server/dispatcher"
"github.com/datacommonsorg/mixer/internal/server/resource"
v2observation "github.com/datacommonsorg/mixer/internal/server/v2/observation"
"github.com/datacommonsorg/mixer/internal/server/v2/resolve"
Expand Down Expand Up @@ -204,45 +207,82 @@ func TestObservationInternal(t *testing.T) {
}
}

type mockObservationDataSource struct {
datasource.DataSource
}

func (m *mockObservationDataSource) Observation(ctx context.Context, req *pbv2.ObservationRequest) (*pbv2.ObservationResponse, error) {
return &pbv2.ObservationResponse{
ByVariable: map[string]*pbv2.VariableObservation{
"Count_Person": {
ByEntity: map[string]*pbv2.EntityObservation{
"country/USA": {},
},
},
},
}, nil
}
func (m *mockObservationDataSource) ID() string { return "mock" }

func TestV2Observation_UsageLog(t *testing.T) {
ctx := metadata.NewIncomingContext(context.Background(), metadata.MD{})
s := &Server{
store: &store.Store{},
metadata: &resource.Metadata{},
flags: &featureflags.Flags{},
writeUsageLogs: true,
}
s.cachedata.Store(&cache.Cache{})
req := &pbv2.ObservationRequest{
Select: []string{"variable", "entity", "date", "value"},
Variable: &pbv2.DcidOrExpression{
Dcids: []string{"Count_Person"},
tests := []struct {
name string
useSpannerGraph bool
}{
{
name: "legacy",
useSpannerGraph: false,
},
Entity: &pbv2.DcidOrExpression{
Dcids: []string{"country/USA"},
{
name: "diverted to dispatcher",
useSpannerGraph: true,
},
}

// Capture slog output
var buf bytes.Buffer
handler := slog.NewTextHandler(&buf, nil)
logger := slog.New(handler)
originalLogger := slog.Default()
slog.SetDefault(logger)
defer slog.SetDefault(originalLogger)
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
ctx := metadata.NewIncomingContext(context.Background(), metadata.MD{})
s := &Server{
store: &store.Store{},
metadata: &resource.Metadata{},
flags: &featureflags.Flags{},
writeUsageLogs: true,
useSpannerGraph: tc.useSpannerGraph,
dispatcher: dispatcher.NewDispatcher(nil, datasources.NewDataSources([]datasource.DataSource{&mockObservationDataSource{}}, nil)),
}
s.cachedata.Store(&cache.Cache{})
req := &pbv2.ObservationRequest{
Select: []string{"variable", "entity", "date", "value"},
Variable: &pbv2.DcidOrExpression{
Dcids: []string{"Count_Person"},
},
Entity: &pbv2.DcidOrExpression{
Dcids: []string{"country/USA"},
},
}

_, _ = s.V2Observation(ctx, req)
// Capture slog output
var buf bytes.Buffer
handler := slog.NewTextHandler(&buf, nil)
logger := slog.New(handler)
originalLogger := slog.Default()
slog.SetDefault(logger)
defer slog.SetDefault(originalLogger)

outStr := strings.TrimSpace(buf.String())
_, _ = s.V2Observation(ctx, req)

// Use regex to match the log message, ignoring the timestamp and pointer address.
wantLogRegex := `time=\S+ level=INFO msg=new_query usage_log.feature="{IsRemote:false Surface:}" usage_log.place_types=\[\] usage_log.query_type=value usage_log.stat_vars=\[0x[0-9a-f]+\] usage_log.response_id=\S+`
matched, err := regexp.MatchString(wantLogRegex, outStr)
if err != nil {
t.Fatalf("Failed to compile regex: %v", err)
}
if !matched {
t.Errorf("log output did not match expected pattern.\nGot: %s\nWant regex: %s", outStr, wantLogRegex)
outStr := strings.TrimSpace(buf.String())

// Use regex to match the log message, ignoring the timestamp and pointer address.
wantLogRegex := `time=\S+ level=INFO msg=new_query usage_log.feature="{IsRemote:false Surface:}" usage_log.place_types=\[\] usage_log.query_type=value usage_log.stat_vars=\[0x[0-9a-f]+\] usage_log.response_id=\S+`
matched, err := regexp.MatchString(wantLogRegex, outStr)
if err != nil {
t.Fatalf("Failed to compile regex: %v", err)
}
if !matched {
t.Errorf("log output did not match expected pattern.\nGot: %s\nWant regex: %s", outStr, wantLogRegex)
}
})
}
}

Expand Down
45 changes: 27 additions & 18 deletions internal/server/remote/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
package remote

import (
"context"
"fmt"
"log/slog"
"net/http"
Expand Down Expand Up @@ -48,14 +49,22 @@ func NewRemoteClient(metadata *resource.Metadata) (*RemoteClient, error) {
}, nil
}

func (rc *RemoteClient) Node(req *pbv2.NodeRequest) (*pbv2.NodeResponse, error) {
func getSurface(ctx context.Context) string {
if ctx == nil {
return ""
}
surface, _ := util.GetMetadata(ctx)
return surface
}

func (rc *RemoteClient) Node(ctx context.Context, req *pbv2.NodeRequest) (*pbv2.NodeResponse, error) {
err := updateNodeRequestNextToken(req, rc.id)
if err != nil {
return nil, err
}

resp := &pbv2.NodeResponse{}
err = util.FetchRemote(rc.metadata, rc.httpClient, "/v2/node", req, resp)
err = util.FetchRemote(rc.metadata, rc.httpClient, "/v2/node", req, resp, getSurface(ctx))
if err != nil {
return nil, err
}
Expand All @@ -68,72 +77,72 @@ func (rc *RemoteClient) Node(req *pbv2.NodeRequest) (*pbv2.NodeResponse, error)
return resp, nil
}

func (rc *RemoteClient) Observation(req *pbv2.ObservationRequest) (*pbv2.ObservationResponse, error) {
func (rc *RemoteClient) Observation(ctx context.Context, req *pbv2.ObservationRequest) (*pbv2.ObservationResponse, error) {
resp := &pbv2.ObservationResponse{}
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/observation", req, resp)
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/observation", req, resp, getSurface(ctx))
if err != nil {
return nil, err
}
return resp, nil
}

func (rc *RemoteClient) NodeSearch(req *pbv2.NodeSearchRequest) (*pbv2.NodeSearchResponse, error) {
func (rc *RemoteClient) NodeSearch(ctx context.Context, req *pbv2.NodeSearchRequest) (*pbv2.NodeSearchResponse, error) {
resp := &pbv2.NodeSearchResponse{}
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v3/node_search", req, resp)
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v3/node_search", req, resp, getSurface(ctx))
if err != nil {
return nil, err
}
return resp, nil
}

func (rc *RemoteClient) Resolve(req *pbv2.ResolveRequest) (*pbv2.ResolveResponse, error) {
func (rc *RemoteClient) Resolve(ctx context.Context, req *pbv2.ResolveRequest) (*pbv2.ResolveResponse, error) {
resp := &pbv2.ResolveResponse{}
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/resolve", req, resp)
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/resolve", req, resp, getSurface(ctx))
if err != nil {
return nil, err
}
return resp, nil
}

func (rc *RemoteClient) Sparql(req *pb.SparqlRequest) (*pb.QueryResponse, error) {
func (rc *RemoteClient) Sparql(ctx context.Context, req *pb.SparqlRequest) (*pb.QueryResponse, error) {
resp := &pb.QueryResponse{}
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/sparql", req, resp)
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/sparql", req, resp, getSurface(ctx))
if err != nil {
return nil, err
}
return resp, nil
}

func (rc *RemoteClient) Event(req *pbv2.EventRequest) (*pbv2.EventResponse, error) {
func (rc *RemoteClient) Event(ctx context.Context, req *pbv2.EventRequest) (*pbv2.EventResponse, error) {
resp := &pbv2.EventResponse{}
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/event", req, resp)
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/event", req, resp, getSurface(ctx))
if err != nil {
return nil, err
}
return resp, nil
}

func (rc *RemoteClient) BulkVariableInfo(req *pbv1.BulkVariableInfoRequest) (*pbv1.BulkVariableInfoResponse, error) {
func (rc *RemoteClient) BulkVariableInfo(ctx context.Context, req *pbv1.BulkVariableInfoRequest) (*pbv1.BulkVariableInfoResponse, error) {
resp := &pbv1.BulkVariableInfoResponse{}
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/bulk/info/variable", req, resp)
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/bulk/info/variable", req, resp, getSurface(ctx))
if err != nil {
return nil, err
}
return resp, nil
}

func (rc *RemoteClient) BulkVariableGroupInfo(req *pbv1.BulkVariableGroupInfoRequest) (*pbv1.BulkVariableGroupInfoResponse, error) {
func (rc *RemoteClient) BulkVariableGroupInfo(ctx context.Context, req *pbv1.BulkVariableGroupInfoRequest) (*pbv1.BulkVariableGroupInfoResponse, error) {
resp := &pbv1.BulkVariableGroupInfoResponse{}
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/bulk/info/variable-group", req, resp)
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/bulk/info/variable-group", req, resp, getSurface(ctx))
if err != nil {
return nil, err
}
return resp, nil
}

func (rc *RemoteClient) FilterStatVarsByEntity(req *pb.FilterStatVarsByEntityRequest) (*pb.FilterStatVarsByEntityResponse, error) {
func (rc *RemoteClient) FilterStatVarsByEntity(ctx context.Context, req *pb.FilterStatVarsByEntityRequest) (*pb.FilterStatVarsByEntityResponse, error) {
resp := &pb.FilterStatVarsByEntityResponse{}
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/variable/filter", req, resp)
err := util.FetchRemote(rc.metadata, rc.httpClient, "/v2/variable/filter", req, resp, getSurface(ctx))
if err != nil {
slog.Error("Failed to fetch remote variable filter", "error", err)
return nil, err
Expand Down
Loading
Loading