Add support for parsing datetime types with type variables

trinodb · Jun 2, 2020 · 8000fab · 8000fab
1 parent 3fc2575
commit 8000fab
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 16 deletions.
diff --git a/presto-parser/src/main/antlr4/io/prestosql/sql/parser/SqlBase.g4 b/presto-parser/src/main/antlr4/io/prestosql/sql/parser/SqlBase.g4
@@ -387,10 +387,10 @@ normalForm
 type
     : ROW '(' rowField (',' rowField)* ')'                                         #rowType
     | INTERVAL from=intervalField (TO to=intervalField)?                           #intervalType
-    | base=TIMESTAMP ('(' precision = INTEGER_VALUE ')')? (WITHOUT TIME ZONE)?     #dateTimeType
-    | base=TIMESTAMP ('(' precision = INTEGER_VALUE ')')? WITH TIME ZONE           #dateTimeType
-    | base=TIME ('(' precision = INTEGER_VALUE ')')? (WITHOUT TIME ZONE)?          #dateTimeType
-    | base=TIME ('(' precision = INTEGER_VALUE ')')? WITH TIME ZONE                #dateTimeType
+    | base=TIMESTAMP ('(' precision = typeParameter ')')? (WITHOUT TIME ZONE)?     #dateTimeType
+    | base=TIMESTAMP ('(' precision = typeParameter ')')? WITH TIME ZONE           #dateTimeType
+    | base=TIME ('(' precision = typeParameter ')')? (WITHOUT TIME ZONE)?          #dateTimeType
+    | base=TIME ('(' precision = typeParameter ')')? WITH TIME ZONE                #dateTimeType
     | DOUBLE PRECISION                                                             #doublePrecisionType
     | ARRAY '<' type '>'                                                           #legacyArrayType
     | MAP '<' keyType=type ',' valueType=type '>'                                  #legacyMapType

diff --git a/presto-parser/src/main/java/io/prestosql/sql/parser/AstBuilder.java b/presto-parser/src/main/java/io/prestosql/sql/parser/AstBuilder.java
@@ -2004,7 +2004,7 @@ else if (context.base.getType() == TIMESTAMP) {
                 getLocation(context),
                 type,
                 context.WITH() != null,
-                getTextIfPresent(context.precision));
+                visitIfPresent(context.precision, DataTypeParameter.class));
     }
 
     @Override

diff --git a/presto-parser/src/main/java/io/prestosql/sql/tree/DateTimeDataType.java b/presto-parser/src/main/java/io/prestosql/sql/tree/DateTimeDataType.java
@@ -31,14 +31,14 @@ public enum Type
 
     private final Type type;
     private final boolean withTimeZone;
-    private final Optional<String> precision;
+    private final Optional<DataTypeParameter> precision;
 
-    public DateTimeDataType(NodeLocation location, Type type, boolean withTimeZone, Optional<String> precision)
+    public DateTimeDataType(NodeLocation location, Type type, boolean withTimeZone, Optional<DataTypeParameter> precision)
     {
         this(Optional.of(location), type, withTimeZone, precision);
     }
 
-    public DateTimeDataType(Optional<NodeLocation> location, Type type, boolean withTimeZone, Optional<String> precision)
+    public DateTimeDataType(Optional<NodeLocation> location, Type type, boolean withTimeZone, Optional<DataTypeParameter> precision)
     {
         super(location);
         this.type = requireNonNull(type, "type is null");
@@ -56,7 +56,7 @@ public boolean isWithTimeZone()
         return withTimeZone;
     }
 
-    public Optional<String> getPrecision()
+    public Optional<DataTypeParameter> getPrecision()
     {
         return precision;
     }

diff --git a/presto-parser/src/test/java/io/prestosql/sql/parser/TestTypeParser.java b/presto-parser/src/test/java/io/prestosql/sql/parser/TestTypeParser.java
@@ -113,13 +113,22 @@ public void testDayTimeTypes()
                 .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, true));
 
         assertThat(type("TIMESTAMP(3)"))
-                .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, false, "3"));
+                .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, false, parameter(location(1, 10), "3")));
 
         assertThat(type("TIMESTAMP(3) WITHOUT TIME ZONE"))
-                .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, false, "3"));
+                .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, false, parameter(location(1, 10), "3")));
 
         assertThat(type("TIMESTAMP(3) WITH TIME ZONE"))
-                .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, true, "3"));
+                .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, true, parameter(location(1, 10), "3")));
+
+        assertThat(type("TIMESTAMP(p)"))
+                .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, false, parameter(simpleType(location(1, 10), "p"))));
+
+        assertThat(type("TIMESTAMP(p) WITHOUT TIME ZONE"))
+                .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, false, parameter(simpleType(location(1, 10), "p"))));
+
+        assertThat(type("TIMESTAMP(p) WITH TIME ZONE"))
+                .isEqualTo(dateTimeType(location(1, 0), TIMESTAMP, true, parameter(simpleType(location(1, 10), "p"))));
 
         assertThat(type("TIME"))
                 .isEqualTo(dateTimeType(location(1, 0), TIME, false));
@@ -131,13 +140,22 @@ public void testDayTimeTypes()
                 .isEqualTo(dateTimeType(location(1, 0), TIME, true));
 
         assertThat(type("TIME(3)"))
-                .isEqualTo(dateTimeType(location(1, 0), TIME, false, "3"));
+                .isEqualTo(dateTimeType(location(1, 0), TIME, false, parameter(location(1, 5), "3")));
 
         assertThat(type("TIME(3) WITHOUT TIME ZONE"))
-                .isEqualTo(dateTimeType(location(1, 0), TIME, false, "3"));
+                .isEqualTo(dateTimeType(location(1, 0), TIME, false, parameter(location(1, 5), "3")));
 
         assertThat(type("TIME(3) WITH TIME ZONE"))
-                .isEqualTo(dateTimeType(location(1, 0), TIME, true, "3"));
+                .isEqualTo(dateTimeType(location(1, 0), TIME, true, parameter(location(1, 5), "3")));
+
+        assertThat(type("TIME(p)"))
+                .isEqualTo(dateTimeType(location(1, 0), TIME, false, parameter(simpleType(location(1, 5), "p"))));
+
+        assertThat(type("TIME(p) WITHOUT TIME ZONE"))
+                .isEqualTo(dateTimeType(location(1, 0), TIME, false, parameter(simpleType(location(1, 5), "p"))));
+
+        assertThat(type("TIME(p) WITH TIME ZONE"))
+                .isEqualTo(dateTimeType(location(1, 0), TIME, true, parameter(simpleType(location(1, 5), "p"))));
     }
 
     @Test

diff --git a/presto-parser/src/test/java/io/prestosql/sql/parser/TreeNodes.java b/presto-parser/src/test/java/io/prestosql/sql/parser/TreeNodes.java
@@ -64,7 +64,7 @@ public static DateTimeDataType dateTimeType(NodeLocation location, DateTimeDataT
         return new DateTimeDataType(location, kind, withTimeZone, Optional.empty());
     }
 
-    public static DateTimeDataType dateTimeType(NodeLocation location, DateTimeDataType.Type kind, boolean withTimeZone, String precision)
+    public static DateTimeDataType dateTimeType(NodeLocation location, DateTimeDataType.Type kind, boolean withTimeZone, DataTypeParameter precision)
     {
         return new DateTimeDataType(location, kind, withTimeZone, Optional.of(precision));
     }