-
Notifications
You must be signed in to change notification settings - Fork 101
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add context propagator and use it to enrich metrics with flow m…
…etadata (#2493) The context is propagated as temporal headers: <img width="1074" alt="image" src="https://github.com/user-attachments/assets/444f8b4d-e909-4b2e-aa3a-03fe0a1ac71e" /> Metrics automatically have the desired attributes: <img width="879" alt="image" src="https://github.com/user-attachments/assets/5dddb50a-52cb-470d-bde2-2bf963deca66" />
- Loading branch information
1 parent
799af39
commit 3664789
Showing
14 changed files
with
288 additions
and
31 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,84 @@ | ||
package shared | ||
|
||
import ( | ||
"context" | ||
|
||
"go.temporal.io/sdk/converter" | ||
"go.temporal.io/sdk/workflow" | ||
|
||
"github.com/PeerDB-io/peerdb/flow/generated/protos" | ||
) | ||
|
||
type TemporalContextKey string | ||
|
||
func (k TemporalContextKey) HeaderKey() string { | ||
return string(k) | ||
} | ||
|
||
const ( | ||
FlowMetadataKey TemporalContextKey = "x-peerdb-flow-metadata" | ||
) | ||
|
||
type PeerMetadata struct { | ||
Name string | ||
Type protos.DBType | ||
} | ||
|
||
func GetFlowMetadata(ctx context.Context) *protos.FlowContextMetadata { | ||
if metadata, ok := ctx.Value(FlowMetadataKey).(*protos.FlowContextMetadata); ok { | ||
return metadata | ||
} | ||
return nil | ||
} | ||
|
||
type ContextPropagator[V any] struct { | ||
Key TemporalContextKey | ||
} | ||
|
||
func NewContextPropagator[V any](key TemporalContextKey) workflow.ContextPropagator { | ||
return &ContextPropagator[V]{Key: key} | ||
} | ||
|
||
func (c *ContextPropagator[V]) Inject(ctx context.Context, writer workflow.HeaderWriter) error { | ||
value := ctx.Value(c.Key) | ||
payload, err := converter.GetDefaultDataConverter().ToPayload(value) | ||
if err != nil { | ||
return err | ||
} | ||
writer.Set(c.Key.HeaderKey(), payload) | ||
return nil | ||
} | ||
|
||
func (c *ContextPropagator[V]) Extract(ctx context.Context, reader workflow.HeaderReader) (context.Context, error) { | ||
if payload, ok := reader.Get(c.Key.HeaderKey()); ok { | ||
var value V | ||
if err := converter.GetDefaultDataConverter().FromPayload(payload, &value); err != nil { | ||
return ctx, nil | ||
} | ||
ctx = context.WithValue(ctx, c.Key, value) | ||
} | ||
|
||
return ctx, nil | ||
} | ||
|
||
func (c *ContextPropagator[V]) InjectFromWorkflow(ctx workflow.Context, writer workflow.HeaderWriter) error { | ||
value := ctx.Value(c.Key) | ||
payload, err := converter.GetDefaultDataConverter().ToPayload(value) | ||
if err != nil { | ||
return err | ||
} | ||
writer.Set(c.Key.HeaderKey(), payload) | ||
return nil | ||
} | ||
|
||
func (c *ContextPropagator[V]) ExtractToWorkflow(ctx workflow.Context, reader workflow.HeaderReader) (workflow.Context, error) { | ||
if payload, ok := reader.Get(c.Key.HeaderKey()); ok { | ||
var value V | ||
if err := converter.GetDefaultDataConverter().FromPayload(payload, &value); err != nil { | ||
return ctx, nil | ||
} | ||
ctx = workflow.WithValue(ctx, c.Key, value) | ||
} | ||
|
||
return ctx, nil | ||
} |
Oops, something went wrong.