From 0a95dba47681c9c4cc4e41ecfb5dadec6357bff6 Mon Sep 17 00:00:00 2001
From: Sri Harsha CH <57220027+harshachinta@users.noreply.github.com>
Date: Thu, 8 Feb 2024 23:53:47 +0530
Subject: [PATCH] fix: cast for Proto type (#2862)

* fix: cast for Proto type

* fix: add null check

* feat(spanner): add ENUM compatibility with getLongArray

* feat: fix argument

* feat: fix bug
---
 .../com/google/cloud/spanner/AbstractResultSet.java   | 11 ++++++++---
 .../google/cloud/spanner/AbstractStructReader.java    |  6 ++++--
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java
index d946257bc45..4904a382f92 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java
@@ -478,7 +478,10 @@ private Object writeReplace() {
           case PROTO:
             builder
                 .set(fieldName)
-                .to(Value.protoMessage((ByteArray) value, fieldType.getProtoTypeFqn()));
+                .to(
+                    Value.protoMessage(
+                        value == null ? null : ((LazyByteArray) value).getByteArray(),
+                        fieldType.getProtoTypeFqn()));
             break;
           case ENUM:
             builder.set(fieldName).to(Value.protoEnum((Long) value, fieldType.getProtoTypeFqn()));
@@ -810,7 +813,8 @@ protected Value getValueInternal(int columnIndex) {
         case INT64:
           return Value.int64(isNull ? null : getLongInternal(columnIndex));
         case ENUM:
-          return Value.protoEnum(getLongInternal(columnIndex), columnType.getProtoTypeFqn());
+          return Value.protoEnum(
+              isNull ? null : getLongInternal(columnIndex), columnType.getProtoTypeFqn());
         case NUMERIC:
           return Value.numeric(isNull ? null : getBigDecimalInternal(columnIndex));
         case PG_NUMERIC:
@@ -826,7 +830,8 @@ protected Value getValueInternal(int columnIndex) {
         case BYTES:
           return Value.internalBytes(isNull ? null : getLazyBytesInternal(columnIndex));
         case PROTO:
-          return Value.protoMessage(getBytesInternal(columnIndex), columnType.getProtoTypeFqn());
+          return Value.protoMessage(
+              isNull ? null : getBytesInternal(columnIndex), columnType.getProtoTypeFqn());
         case TIMESTAMP:
           return Value.timestamp(isNull ? null : getTimestampInternal(columnIndex));
         case DATE:
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractStructReader.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractStructReader.java
index 53c49a5a54b..ef6f63d52ea 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractStructReader.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractStructReader.java
@@ -340,14 +340,16 @@ public List<Boolean> getBooleanList(String columnName) {
 
   @Override
   public long[] getLongArray(int columnIndex) {
-    checkNonNullOfType(columnIndex, Type.array(Type.int64()), columnIndex);
+    checkNonNullOfCodes(columnIndex, Collections.singletonList(Code.ARRAY), columnIndex);
+    checkArrayElementType(columnIndex, Arrays.asList(Code.ENUM, Code.INT64), columnIndex);
     return getLongArrayInternal(columnIndex);
   }
 
   @Override
   public long[] getLongArray(String columnName) {
     int columnIndex = getColumnIndex(columnName);
-    checkNonNullOfType(columnIndex, Type.array(Type.int64()), columnName);
+    checkNonNullOfCodes(columnIndex, Collections.singletonList(Code.ARRAY), columnName);
+    checkArrayElementType(columnIndex, Arrays.asList(Code.ENUM, Code.INT64), columnName);
     return getLongArrayInternal(columnIndex);
   }