Merge pull request #84 from GoogleCloudPlatform/main

Sync main
ollionorg · Jan 28, 2025 · f1a4d1e · f1a4d1e
2 parents 5d5f1a8 + 208e2ba
commit f1a4d1e
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 11 deletions.
diff --git a/...rc/main/java/com/google/cloud/teleport/v2/templates/dbutils/dml/CassandraTypeHandler.java b/...rc/main/java/com/google/cloud/teleport/v2/templates/dbutils/dml/CassandraTypeHandler.java
@@ -118,8 +118,11 @@ private interface HandlerSupplier<T> {
    * @throws IllegalArgumentException If the value is neither a valid number string, byte array, nor
    *     a valid {@link ByteBuffer} for varint representation.
    */
-  private static BigInteger handleCassandraVarintType(String value) {
-    return new BigInteger(value);
+  private static BigInteger handleCassandraVarintType(Object value) {
+    if (value instanceof byte[]) {
+      return new BigInteger((byte[]) value);
+    }
+    return new BigInteger(value.toString());
   }
 
   /**
@@ -353,6 +356,10 @@ private static Object handleSpannerColumnType(
   private static PreparedStatementValueObject<?> parseAndCastToCassandraType(
       String columnType, Object colValue) {
 
+    if (columnType.startsWith("frozen<")) {
+      return parseAndCastToCassandraType(extractInnerType(columnType), colValue);
+    }
+
     // Handle collection types
     if (columnType.startsWith("list<")) {
       return safeHandle(
@@ -443,7 +450,7 @@ private static PreparedStatementValueObject<?> parseAndCastToCassandraType(
 
       case "varint":
         return PreparedStatementValueObject.create(
-            columnType, safeHandle(() -> handleCassandraVarintType(colValue.toString())));
+            columnType, safeHandle(() -> handleCassandraVarintType(colValue)));
 
       case "duration":
         return PreparedStatementValueObject.create(
@@ -523,9 +530,10 @@ private static BigDecimal parseDecimal(Object colValue) {
    * @throws IllegalArgumentException if the column type is invalid or the value cannot be parsed.
    */
   private static Object parseFloatingPoint(String columnType, Object colValue) {
-    return columnType.equals("double")
-        ? Double.parseDouble((String) colValue)
-        : Float.parseFloat((String) colValue);
+    if (columnType.equals("double")) {
+      return Double.parseDouble((String) colValue);
+    }
+    return Float.parseFloat((String) colValue);
   }
 
   private static LocalDate parseDate(Object colValue) {

diff --git a/...rcedb/src/main/java/com/google/cloud/teleport/v2/templates/transforms/SourceWriterFn.java b/...rcedb/src/main/java/com/google/cloud/teleport/v2/templates/transforms/SourceWriterFn.java
@@ -15,6 +15,8 @@
  */
 package com.google.cloud.teleport.v2.templates.transforms;
 
+import com.datastax.oss.driver.api.core.servererrors.QueryExecutionException;
+import com.datastax.oss.driver.api.core.type.codec.CodecNotFoundException;
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -233,7 +235,9 @@ public void processElement(ProcessContext c) {
       } catch (InvalidTransformationException ex) {
         invalidTransformationException.inc();
         outputWithTag(c, Constants.PERMANENT_ERROR_TAG, ex.getMessage(), spannerRec);
-      } catch (ChangeEventConvertorException ex) {
+      } catch (ChangeEventConvertorException
+          | CodecNotFoundException
+          | QueryExecutionException ex) {
         outputWithTag(c, Constants.PERMANENT_ERROR_TAG, ex.getMessage(), spannerRec);
       } catch (SpannerException
           | IllegalStateException

diff --git a/...est/java/com/google/cloud/teleport/v2/templates/dbutils/dml/CassandraTypeHandlerTest.java b/...est/java/com/google/cloud/teleport/v2/templates/dbutils/dml/CassandraTypeHandlerTest.java
@@ -769,7 +769,7 @@ public void testGetColumnValueByTypeForFloat() {
     String columnName = "test_column";
 
     JSONObject valuesJson = new JSONObject();
-    valuesJson.put(columnName, new BigDecimal("5.5"));
+    valuesJson.put(columnName, 5.5f);
 
     SpannerColumnDefinition spannerColDef = new SpannerColumnDefinition(columnName, spannerType);
     SourceColumnDefinition sourceColDef = new SourceColumnDefinition(columnName, sourceColType);
@@ -779,7 +779,29 @@ public void testGetColumnValueByTypeForFloat() {
     assertTrue(result instanceof PreparedStatementValueObject);
 
     Object actualValue = ((PreparedStatementValueObject<?>) result).value();
-    assertEquals(new BigDecimal(5.5), actualValue);
+    assertEquals(5.5f, actualValue);
+  }
+
+  @Test
+  public void testGetColumnValueByTypeForFloatIllegalArgumentException() {
+    SpannerColumnType spannerType = new SpannerColumnType("date", false);
+    SourceColumnType sourceColType = new SourceColumnType("float", null, null);
+    String columnName = "test_column";
+
+    JSONObject valuesJson = new JSONObject();
+    valuesJson.put(columnName, "2024-12-12");
+
+    SpannerColumnDefinition spannerColDef = new SpannerColumnDefinition(columnName, spannerType);
+    SourceColumnDefinition sourceColDef = new SourceColumnDefinition(columnName, sourceColType);
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          PreparedStatementValueObject<?> preparedStatementValueObject =
+              getColumnValueByType(spannerColDef, sourceColDef, valuesJson, null);
+          castToExpectedType(
+              preparedStatementValueObject.dataType(), preparedStatementValueObject.value());
+        });
   }
 
   @Test
@@ -911,7 +933,7 @@ public void testGetColumnValueByTypeForFloatFromString() {
 
     Object castResult = CassandraTypeHandler.castToExpectedType(result.dataType(), result.value());
 
-    assertEquals(5.5, castResult);
+    assertEquals(5.5f, castResult);
   }
 
   @Test
@@ -1036,7 +1058,7 @@ public void testCastToExpectedTypeForVariousTypes() throws UnknownHostException
         new BigDecimal("123.456"),
         castToExpectedType("decimal", new BigDecimal("123.456").toString()));
     assertEquals(123.456, castToExpectedType("double", "123.456"));
-    assertEquals(123.45f, ((Double) castToExpectedType("float", "123.45")).floatValue(), 0.00001);
+    assertEquals(123.45f, castToExpectedType("float", "123.45"));
     assertEquals(InetAddress.getByName("127.0.0.1"), castToExpectedType("inet", "127.0.0.1"));
     assertEquals(123, castToExpectedType("int", "123"));
     assertEquals((short) 123, castToExpectedType("smallint", "123"));
@@ -1074,6 +1096,18 @@ public void testCastToExpectedTypeForJSONObjectStringifyToMap() {
         });
   }
 
+  @Test
+  public void testCastToExpectedTypeForJSONObjectStringifyToMFrozenMap() {
+    String cassandraType = "frozen<map<int, text>>";
+    String columnValue = "{\"1\": \"One\", \"2\": \"Two\"}";
+    Object castResult = castToExpectedType(cassandraType, columnValue);
+    assertTrue(castResult instanceof Map);
+    assertTrue(((Map<?, ?>) castResult).containsKey(1));
+    assertTrue(((Map<?, ?>) castResult).containsKey(2));
+    assertEquals("One", ((Map<?, ?>) castResult).get(1));
+    assertEquals("Two", ((Map<?, ?>) castResult).get(2));
+  }
+
   @Test
   public void testCastToExpectedTypeForExceptionScenario() {
     String cassandraType = "int";
@@ -1181,6 +1215,21 @@ public void testHandleCassandraVarintType_String() {
     assertEquals(expected, result);
   }
 
+  @Test
+  public void testHandleCassandraVarintType_ForBytesArray() {
+    byte[] byteArray = new byte[] {0, 0, 0, 0, 0, 0, 0, 10};
+    BigInteger expected = new BigInteger(byteArray);
+    Object result = CassandraTypeHandler.castToExpectedType("varint", byteArray);
+    assertEquals(expected, result);
+  }
+
+  @Test
+  public void testHandleCassandraVarintType_ForInteger() {
+    Long inputValue = 123456789L;
+    Object result = CassandraTypeHandler.castToExpectedType("varint", inputValue);
+    assertEquals(BigInteger.valueOf(inputValue), result);
+  }
+
   @Test
   public void testHandleCassandraVarintType_InvalidString() {
     String invalidString = "invalid-number";