Merge pull request #1442 from cloudsufi/patch/time-partion-col-validation

[PLUGIN-1805] Fix TimePartitioningColumn DATETIME validation

Author: psainics

docs/BigQueryTable-batchsink.md

@@ -110,6 +110,17 @@ is ignored if the table already exists.
 **Time Partitioning Type**: Specifies the time partitioning type. Can either be Daily or Hourly or Monthly or Yearly.
 Default is Daily. Ignored when table already exists
 
+> The table below shows the compatibility of different time schema types with various time partitioning types in BigQuery.
+
+| Schema Type / Partion Type  | Hourly  | Daily   | Monthly | Yearly  |
+|-------------------------| ------- | ------- | ------- | ------- |
+| TIMESTAMP_MILLIS        | ✓ | ✓ | ✓ | ✓ |
+| TIMESTAMP_MICROS        | ✓ | ✓ | ✓ | ✓ |
+| DATETIME                | ✓ | ✓ | ✓ | ✓ |
+| DATE                    | ✗ | ✓ | ✓ | ✓ |
+| TIME_MILLIS             | ✗ | ✗ | ✗ | ✗ |
+| TIME_MICROS             | ✗ | ✗ | ✗ | ✗ |
+
 **Range Start**: For integer partitioning, specifies the start of the range. Only used when table doesn’t 
 exist already, and partitioning type is set to Integer.
 * The start value is inclusive.

src/main/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySinkConfig.java

@@ -526,23 +526,21 @@ private void validateTimePartitioningColumn(String columnName, FailureCollector
 
     boolean isTimestamp = logicalType == LogicalType.TIMESTAMP_MICROS || logicalType == LogicalType.TIMESTAMP_MILLIS;
     boolean isDate = logicalType == LogicalType.DATE;
-    boolean isTimestampOrDate = isTimestamp || isDate;
-
-    // If timePartitioningType is HOUR, then logicalType cannot be DATE Only TIMESTAMP_MICROS and TIMESTAMP_MILLIS
-    if (timePartitioningType == TimePartitioning.Type.HOUR && !isTimestamp) {
-      collector.addFailure(
-                      String.format("Partition column '%s' is of invalid type '%s'.",
+    boolean isDateTime = logicalType == LogicalType.DATETIME;
+    boolean isTimestampOrDateOrDateTime = isTimestamp || isDate || isDateTime;
+    boolean isTimestampOrDateTime = isTimestamp || isDateTime;
+    // TimePartitioningType HOUR is supported by TIMESTAMP_MICROS, TIMESTAMP_MILLIS, DATETIME
+    if (timePartitioningType == TimePartitioning.Type.HOUR && !isTimestampOrDateTime) {
+      collector.addFailure(String.format("Partition column '%s' is of invalid type '%s'.",
                               columnName, fieldSchema.getDisplayName()),
-                      "Partition column must be a timestamp.").withConfigProperty(NAME_PARTITION_BY_FIELD)
-              .withOutputSchemaField(columnName).withInputSchemaField(columnName);
-
-    // For any other timePartitioningType (DAY, MONTH, YEAR) logicalType can be DATE, TIMESTAMP_MICROS, TIMESTAMP_MILLIS
-    } else if (!isTimestampOrDate) {
-      collector.addFailure(
-                      String.format("Partition column '%s' is of invalid type '%s'.",
+                      "Partition column must be of type TIMESTAMP or DATETIME")
+        .withConfigProperty(NAME_PARTITION_BY_FIELD).withOutputSchemaField(columnName).withInputSchemaField(columnName);
+    // TimePartitioningType (DAY, MONTH, YEAR) are supported by TIMESTAMP_MICROS, TIMESTAMP_MILLIS, DATE, DATETIME
+    } else if (!isTimestampOrDateOrDateTime) {
+      collector.addFailure(String.format("Partition column '%s' is of invalid type '%s'.",
                               columnName, fieldSchema.getDisplayName()),
-                      "Partition column must be a date or timestamp.").withConfigProperty(NAME_PARTITION_BY_FIELD)
-              .withOutputSchemaField(columnName).withInputSchemaField(columnName);
+                      "Partition column must be of type TIMESTAMP, DATE or DATETIME")
+        .withConfigProperty(NAME_PARTITION_BY_FIELD).withOutputSchemaField(columnName).withInputSchemaField(columnName);
     }
   }
 

src/test/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySinkConfigTest.java

@@ -129,6 +129,34 @@ public void testValidateTimePartitioningColumnWithNullAndDate() throws
     Assert.assertEquals(0, collector.getValidationFailures().size());
   }
 
+  @Test
+  public void testValidateTimePartitioningColumnWithMonthAndDateTime() throws
+    InvocationTargetException, IllegalAccessException {
+
+    String columnName = "partitionFrom";
+    Schema schema = Schema.of(Schema.LogicalType.DATETIME);
+
+    Schema fieldSchema = schema.isNullable() ? schema.getNonNullable() : schema;
+    TimePartitioning.Type timePartitioningType = TimePartitioning.Type.MONTH;
+
+    validateTimePartitioningColumnMethod.invoke(config, columnName, collector, fieldSchema, timePartitioningType);
+    Assert.assertEquals(0, collector.getValidationFailures().size());
+  }
+
+  @Test
+  public void testValidateTimePartitioningColumnWithHourAndDateTime() throws
+    InvocationTargetException, IllegalAccessException {
+
+    String columnName = "partitionFrom";
+    Schema schema = Schema.of(Schema.LogicalType.DATETIME);
+
+    Schema fieldSchema = schema.isNullable() ? schema.getNonNullable() : schema;
+    TimePartitioning.Type timePartitioningType = TimePartitioning.Type.HOUR;
+
+    validateTimePartitioningColumnMethod.invoke(config, columnName, collector, fieldSchema, timePartitioningType);
+    Assert.assertEquals(0, collector.getValidationFailures().size());
+  }
+
   @Test
   public void testValidateColumnNameWithValidColumnName() {
     String columnName = "test";