From 77b62fd409669f7622d06a992fc78d9a225b5a35 Mon Sep 17 00:00:00 2001
From: Raj Nishtala <rnishtala@sumologic.com>
Date: Wed, 8 May 2024 12:32:11 -0400
Subject: [PATCH] fix(ottl): Replace keys in a map after applying the optional
 function on it

---
 .chloggen/replace_pattern_key.yaml            | 27 +++++++++++++++++++
 .../ottlfuncs/func_replace_all_patterns.go    | 13 ++++++++-
 .../func_replace_all_patterns_test.go         | 21 +++++++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 .chloggen/replace_pattern_key.yaml

diff --git a/.chloggen/replace_pattern_key.yaml b/.chloggen/replace_pattern_key.yaml
new file mode 100644
index 000000000000..e4d72ddf695d
--- /dev/null
+++ b/.chloggen/replace_pattern_key.yaml
@@ -0,0 +1,27 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: bug_fix
+
+# The name of the component, or a single word describing the area of concern, (e.g. filelogreceiver)
+component: ottl
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: Replace keys in a map after applying the optional function on it.
+
+# Mandatory: One or more tracking issues related to the change. You can use the PR number here if no issue exists.
+issues: [32896]
+
+# (Optional) One or more lines of additional information to render under the primary note.
+# These lines will be padded with 2 spaces and then inserted directly into the document.
+# Use pipe (|) for multiline entries.
+subtext:
+
+# If your change doesn't affect end users or the exported elements of any package,
+# you should instead start your pull request title with [chore] or use the "Skip Changelog" label.
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if there is a change to a library API.
+# Default: '[user]'
+change_logs: []
diff --git a/pkg/ottl/ottlfuncs/func_replace_all_patterns.go b/pkg/ottl/ottlfuncs/func_replace_all_patterns.go
index 0f3aae69e8c2..d6081fe068ff 100644
--- a/pkg/ottl/ottlfuncs/func_replace_all_patterns.go
+++ b/pkg/ottl/ottlfuncs/func_replace_all_patterns.go
@@ -86,7 +86,18 @@ func replaceAllPatterns[K any](target ottl.PMapGetter[K], mode string, regexPatt
 						if err != nil {
 							return false
 						}
-						updated.PutStr(key, updatedString)
+						switch originalValue.Type() {
+						case pcommon.ValueTypeStr:
+							updated.PutStr(updatedString, originalValue.Str())
+						case pcommon.ValueTypeInt:
+							updated.PutInt(updatedString, originalValue.Int())
+						case pcommon.ValueTypeDouble:
+							updated.PutDouble(updatedString, originalValue.Double())
+						case pcommon.ValueTypeBool:
+							updated.PutBool(updatedString, originalValue.Bool())
+						default:
+							// handle other types here
+						}
 					} else {
 						updatedKey := compiledPattern.ReplaceAllString(key, replacementVal)
 						originalValue.CopyTo(updated.PutEmpty(updatedKey))
diff --git a/pkg/ottl/ottlfuncs/func_replace_all_patterns_test.go b/pkg/ottl/ottlfuncs/func_replace_all_patterns_test.go
index 534bd80f2c07..093aa09e1577 100644
--- a/pkg/ottl/ottlfuncs/func_replace_all_patterns_test.go
+++ b/pkg/ottl/ottlfuncs/func_replace_all_patterns_test.go
@@ -79,6 +79,27 @@ func Test_replaceAllPatterns(t *testing.T) {
 				expectedMap.PutBool("test6", true)
 			},
 		},
+		{
+			name:    "replace only matches (with hash function)",
+			target:  target,
+			mode:    modeKey,
+			pattern: "test",
+			replacement: ottl.StandardStringGetter[pcommon.Map]{
+				Getter: func(context.Context, pcommon.Map) (any, error) {
+					return "test1", nil
+				},
+			},
+			replacementFormat: ottl.Optional[ottl.StringGetter[pcommon.Map]]{},
+			function:          optionalArg,
+			want: func(expectedMap pcommon.Map) {
+				expectedMap.PutStr("hash(test1)", "hello world")
+				expectedMap.PutStr("hash(test1)2", "hello")
+				expectedMap.PutStr("hash(test1)3", "goodbye world1 and world2")
+				expectedMap.PutInt("hash(test1)4", 1234)
+				expectedMap.PutDouble("hash(test1)5", 1234)
+				expectedMap.PutBool("hash(test1)6", true)
+			},
+		},
 		{
 			name:    "replace only matches (with capture group and hash function)",
 			target:  target,