From c683a8c2049eae95df86a293b9b8be353741a675 Mon Sep 17 00:00:00 2001
From: Saggi Mizrahi <smizrahi@google.com>
Date: Sun, 29 Oct 2023 09:17:20 +0000
Subject: [PATCH] Add support for protbuf enums

---
 tools/mc2bq/pkg/schema/schema.go                     | 4 ++++
 tools/mc2bq/pkg/schema/schema_test.go                | 4 ++++
 tools/mc2bq/pkg/schema/testdata/exporter.json.golden | 1 +
 3 files changed, 9 insertions(+)

diff --git a/tools/mc2bq/pkg/schema/schema.go b/tools/mc2bq/pkg/schema/schema.go
index 1f63af4..dfd5709 100644
--- a/tools/mc2bq/pkg/schema/schema.go
+++ b/tools/mc2bq/pkg/schema/schema.go
@@ -28,6 +28,8 @@ import (
 	"time"
 
 	"cloud.google.com/go/bigquery"
+	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/GoogleCloudPlatform/migrationcenter-utils/tools/mc2bq/pkg/messages"
@@ -230,6 +232,8 @@ func normalizeToSchema(obj any, schema *bigquery.FieldSchema) (any, error) {
 		switch obj := objValue.Interface().(type) {
 		case string:
 			return obj, nil
+		case protoreflect.Enum:
+			return protoimpl.X.EnumStringOf(obj.Descriptor(), protoreflect.EnumNumber(obj.Number())), nil
 		default:
 			return nil, wrapWithSerializeError(schema.Name, fmt.Errorf("convert field of type %q to string", reflect.TypeOf(obj).String()))
 		}
diff --git a/tools/mc2bq/pkg/schema/schema_test.go b/tools/mc2bq/pkg/schema/schema_test.go
index 8c97d27..03fafc8 100644
--- a/tools/mc2bq/pkg/schema/schema_test.go
+++ b/tools/mc2bq/pkg/schema/schema_test.go
@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/bigquery"
+	"cloud.google.com/go/migrationcenter/apiv1/migrationcenterpb"
 	"github.com/google/go-cmp/cmp"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
@@ -55,6 +56,7 @@ func TestObjectSerializer(t *testing.T) {
 		RecordArray     []*TestInnerStruct
 		RecordMap       map[string]*TestInnerStruct
 		DeprecatedField string
+		Enum            migrationcenterpb.CommitmentPlan
 	}
 	schema := bigquery.Schema{
 		{Name: "int_scalar", Type: bigquery.IntegerFieldType},
@@ -76,6 +78,7 @@ func TestObjectSerializer(t *testing.T) {
 			{Name: "key", Type: bigquery.StringFieldType},
 			{Name: "value", Type: bigquery.RecordFieldType, Schema: innerStructSchema},
 		}},
+		{Name: "enum", Type: bigquery.StringFieldType},
 		// DeprecatedField doesn't appear in the schema as it has been deprecated
 	}
 	// This ensures that we are testing all the field types that exist in the schema
@@ -97,6 +100,7 @@ func TestObjectSerializer(t *testing.T) {
 		RecordArray:     []*TestInnerStruct{{FieldA: "1", FieldB: 1}, {FieldA: "2", FieldB: 2}},
 		RecordMap:       map[string]*TestInnerStruct{"key": {FieldA: "value"}},
 		DeprecatedField: "I should not be in the result JSON",
+		Enum:            1,
 	}
 
 	got, err := SerializeObjectToBigQuery(obj, "object", schema)
diff --git a/tools/mc2bq/pkg/schema/testdata/exporter.json.golden b/tools/mc2bq/pkg/schema/testdata/exporter.json.golden
index 05b9f8a..e4f2120 100644
--- a/tools/mc2bq/pkg/schema/testdata/exporter.json.golden
+++ b/tools/mc2bq/pkg/schema/testdata/exporter.json.golden
@@ -10,6 +10,7 @@
       "value": "bar"
     }
   ],
+  "enum": "COMMITMENT_PLAN_NONE",
   "float32_scalar": 0.32,
   "float64_scalar": 0.64,
   "int_array": [