From 6dcecc4d71a01542c10b1f0ad8bf07a6656a4ec7 Mon Sep 17 00:00:00 2001
From: Liran Mimony <liran@odigos.io>
Date: Mon, 25 May 2026 11:49:06 +0300
Subject: [PATCH] chore: New Simple Demo apps feature: Push special HTTP header
 for better instrumentation rules testing

---
 inventory/inventory/main.py | 10 +++++++---
 membership/main.go          | 22 +++++++++++++++++-----
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/inventory/inventory/main.py b/inventory/inventory/main.py
index fca6410..46705a7 100644
--- a/inventory/inventory/main.py
+++ b/inventory/inventory/main.py
@@ -6,7 +6,7 @@
 dep_location = os.path.normpath(os.path.join(os.path.dirname(__file__), '..', 'site-packages'))
 sys.path.append(dep_location)
 
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response
 import time
 import signal
 
@@ -47,15 +47,19 @@ def __init__(self, id, name, image):
 
 @app.route('/inventory', methods=['GET'])
 def get_inventory():
+    resp = Response(jsonify([item.__dict__ for item in inventory_items]))
+    resp.headers.add("odigos-special-header", "inventory-service")
     logging.info("Returning inventory")
-    return jsonify([item.__dict__ for item in inventory_items])
+    return resp
 
 @app.route('/buy', methods=['POST'])
 def buy_product():
     product_id = request.args.get('id', type=int)
+    resp = Response(jsonify({"message": "Product purchased successfully"}))
+    resp.headers.add("odigos-special-header", "inventory-service")
     logging.info(f"Buying product with id {product_id}")
     time.sleep(1)
-    return jsonify({"message": "Product purchased successfully"})
+    return resp
 
 def signal_handler(sig, frame):
     logging.info('Terminating inventory service')
diff --git a/membership/main.go b/membership/main.go
index e314735..cb3270a 100644
--- a/membership/main.go
+++ b/membership/main.go
@@ -15,7 +15,7 @@ import (
 )
 
 const (
-	httpPort = 8080
+	httpPort               = 8080
 	pricingEndpointEnvName = "PRICING_SERVICE_HOST"
 )
 
@@ -23,6 +23,12 @@ type IsMemberResponse struct {
 	IsMember bool `json:"isMember"`
 }
 
+func (imr *IsMemberResponse) CheckIsMember() bool {
+	return imr.IsMember
+}
+
+var tracer = otel.Tracer("github.com/keyval-dev/test-apps/kv-shop/membership")
+
 func getPortFromEnvOrDefault() int {
 	port := httpPort
 	if envPort := os.Getenv("PORT"); envPort != "" {
@@ -36,8 +42,6 @@ func getPortFromEnvOrDefault() int {
 	return port
 }
 
-var tracer = otel.Tracer("github.com/keyval-dev/test-apps/kv-shop/membership")
-
 func checkIfMember(ctx context.Context, userId uuid.UUID) bool {
 	_, span := tracer.Start(ctx, "checkIfMember")
 	defer span.End()
@@ -45,7 +49,10 @@ func checkIfMember(ctx context.Context, userId uuid.UUID) bool {
 	slog.Info("Checking if user is a member", "userId", userId)
 
 	// arbitrary decide if user is a member by it's id
-	isMember := userId.String()[29] >= '0' && userId.String()[29] <= '6'
+	res := IsMemberResponse{
+		IsMember: userId.String()[29] >= '0' && userId.String()[29] <= '6',
+	}
+	isMember := res.CheckIsMember()
 	span.SetAttributes(attribute.String("simple_demo.userId", userId.String()), attribute.Bool("simple_demo.is_member", isMember))
 
 	return isMember
@@ -54,9 +61,14 @@ func checkIfMember(ctx context.Context, userId uuid.UUID) bool {
 func main() {
 	port := getPortFromEnvOrDefault()
 	slog.Info("Starting Membership service", "port", port)
+
 	http.HandleFunc("/isMember", func(writer http.ResponseWriter, request *http.Request) {
 		slog.Info("isMember called")
 
+		// inject odigos-special-header
+		writer.Header().Set("odigos-special-header", "membership-service")
+		request.Header.Set("odigos-special-header", "membership-service")
+
 		userId := uuid.New()
 
 		isMember := checkIfMember(request.Context(), userId)
@@ -77,7 +89,7 @@ func main() {
 			writer.WriteHeader(http.StatusInternalServerError)
 			return
 		}
-
+		req.Header.Set("odigos-special-header", "membership-service")
 		_, err = http.DefaultClient.Do(req)
 		if err != nil {
 			slog.Error("failed to call pricing service", "error", err)