From 083c889c70a2fd2fae62b92aefeaaa73a2706c34 Mon Sep 17 00:00:00 2001
From: Elia Bracci <ebracci@sumologic.com>
Date: Fri, 20 Oct 2023 16:32:34 +0200
Subject: [PATCH] Added method OriginalString

---
 protocol.go      | 25 +++++++++++++++++--------
 protocol_test.go | 11 ++++++++++-
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/protocol.go b/protocol.go
index 28608eb8..fdc9b8ed 100644
--- a/protocol.go
+++ b/protocol.go
@@ -49,6 +49,8 @@ type Protocol struct {
 	doc string
 
 	hash string
+
+	originalString string
 }
 
 // NewProtocol creates a protocol instance.
@@ -56,6 +58,7 @@ func NewProtocol(
 	name, namepsace string,
 	types []NamedSchema,
 	messages map[string]*Message,
+	protocolString string,
 	opts ...ProtocolOption,
 ) (*Protocol, error) {
 	var cfg protocolConfig
@@ -69,11 +72,12 @@ func NewProtocol(
 	}
 
 	p := &Protocol{
-		name:       n,
-		properties: newProperties(cfg.props, protocolReserved),
-		types:      types,
-		messages:   messages,
-		doc:        cfg.doc,
+		name:           n,
+		properties:     newProperties(cfg.props, protocolReserved),
+		types:          types,
+		messages:       messages,
+		doc:            cfg.doc,
+		originalString: protocolString,
 	}
 
 	b := md5.Sum([]byte(p.String()))
@@ -102,6 +106,11 @@ func (p *Protocol) Types() []NamedSchema {
 	return p.types
 }
 
+// OriginalString returns the original string of the protocol.
+func (p *Protocol) OriginalString() string {
+	return p.originalString
+}
+
 // String returns the canonical form of the protocol.
 func (p *Protocol) String() string {
 	types := ""
@@ -231,7 +240,7 @@ func ParseProtocol(protocol string) (*Protocol, error) {
 	}
 
 	seen := seenCache{}
-	return parseProtocol(m, seen, cache)
+	return parseProtocol(m, seen, cache, protocol)
 }
 
 type protocol struct {
@@ -243,7 +252,7 @@ type protocol struct {
 	Props     map[string]any            `mapstructure:",remain"`
 }
 
-func parseProtocol(m map[string]any, seen seenCache, cache *SchemaCache) (*Protocol, error) {
+func parseProtocol(m map[string]any, seen seenCache, cache *SchemaCache, protocolString string) (*Protocol, error) {
 	var (
 		p    protocol
 		meta mapstructure.Metadata
@@ -279,7 +288,7 @@ func parseProtocol(m map[string]any, seen seenCache, cache *SchemaCache) (*Proto
 		}
 	}
 
-	return NewProtocol(p.Protocol, p.Namespace, types, messages, WithProtoDoc(p.Doc), WithProtoProps(p.Props))
+	return NewProtocol(p.Protocol, p.Namespace, types, messages, protocolString, WithProtoDoc(p.Doc), WithProtoProps(p.Props))
 }
 
 func parseProtocolTypes(namespace string, types []any, seen seenCache, cache *SchemaCache) ([]NamedSchema, error) {
diff --git a/protocol_test.go b/protocol_test.go
index 4eb2b970..c919311a 100644
--- a/protocol_test.go
+++ b/protocol_test.go
@@ -21,11 +21,20 @@ func TestMustParseProtocol_PanicsOnError(t *testing.T) {
 }
 
 func TestNewProtocol_ValidatesName(t *testing.T) {
-	_, err := avro.NewProtocol("0test", "", nil, nil)
+	_, err := avro.NewProtocol("0test", "", nil, nil, "")
 
 	assert.Error(t, err)
 }
 
+func TestNewProtocol_OriginalString(t *testing.T) {
+	schema := `{"protocol":"test", "namespace": "org.hamba.avro", "messages":{"test":{"request": [{"name": "foobar", "type": "string"}]}}}`
+
+	proto, err := avro.ParseProtocol(schema)
+
+	assert.NoError(t, err)
+	assert.Equal(t, schema, proto.OriginalString())
+}
+
 func TestNewMessage(t *testing.T) {
 	field, _ := avro.NewField("test", avro.NewPrimitiveSchema(avro.String, nil))
 	fields := []*avro.Field{field}