sql (fix): Support TIMESTAMP AT TIME ZONE literal (#3576)

wvlet · Jul 2, 2024 · b123a75 · b123a75
1 parent 9fa19d0
commit b123a75
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 2 deletions.
diff --git a/airframe-sql/src/main/scala/wvlet/airframe/sql/model/Expression.scala b/airframe-sql/src/main/scala/wvlet/airframe/sql/model/Expression.scala
@@ -1157,6 +1157,14 @@ object Expression {
     override def sqlExpr             = s"TIMESTAMP '${value}'"
     override def toString            = s"Literal(TIMESTAMP '${value}')"
   }
+  case class TimestampWithTimeZoneLiteral(value: String, timezone: String, nodeLocation: Option[NodeLocation])
+      extends Literal
+      with LeafExpression {
+    override def dataType: DataType  = DataType.TimestampType(TimestampField.TIMESTAMP, false)
+    override def stringValue: String = value
+    override def sqlExpr             = s"TIMESTAMP '${value}' AT TIME ZONE '${timezone}'"
+    override def toString            = s"Literal(TIMESTAMP '${value}' AT '${timezone}')"
+  }
   case class DecimalLiteral(value: String, nodeLocation: Option[NodeLocation]) extends Literal with LeafExpression {
     override def dataType: DataType  = DataType.DecimalType(TypeVariable("precision"), TypeVariable("scale"))
     override def stringValue: String = value

diff --git a/airframe-sql/src/main/scala/wvlet/airframe/sql/parser/SQLInterpreter.scala b/airframe-sql/src/main/scala/wvlet/airframe/sql/parser/SQLInterpreter.scala
@@ -471,12 +471,20 @@ class SQLInterpreter(withNodeLocation: Boolean = true) extends SqlBaseBaseVisito
         // TODO Parse decimal-type precision properly
         case "decimal" => DecimalLiteral(v, getLocation(ctx))
         case "char"    => CharLiteral(v, getLocation(ctx))
-        case other =>
-          GenericLiteral(tpe, v, getLocation(ctx))
+        case other     => GenericLiteral(tpe, v, getLocation(ctx))
       }
     }
   }
 
+  override def visitAtTimeZone(ctx: AtTimeZoneContext): Expression = {
+    val v = expression(ctx.timeZoneSpecifier()).asInstanceOf[StringLiteral].value
+
+    expression(ctx.valueExpression()) match {
+      case t: TimestampLiteral => TimestampWithTimeZoneLiteral(t.value, v, t.nodeLocation)
+      case other               => other
+    }
+  }
+
   override def visitBasicStringLiteral(ctx: BasicStringLiteralContext): StringLiteral = {
     StringLiteral(unquote(ctx.STRING().getText), getLocation(ctx))
   }

diff --git a/airframe-sql/src/test/scala/wvlet/airframe/sql/parser/SQLGeneratorTest.scala b/airframe-sql/src/test/scala/wvlet/airframe/sql/parser/SQLGeneratorTest.scala
@@ -186,4 +186,12 @@ class SQLGeneratorTest extends AirSpec {
     val sql = SQLGenerator.print(resolvedPlan).toLowerCase
     sql shouldBe "select * from (select * from a) join a using (id)"
   }
+
+  test("print TIMESTAMP AT TIME ZONE") {
+    val resolvedPlan =
+      SQLAnalyzer.analyze("SELECT TIMESTAMP '1992-02-01 00:00 UTC' AT TIME ZONE 'Asia/Tokyo'", "default", demoCatalog)
+
+    val sql = SQLGenerator.print(resolvedPlan)
+    sql shouldBe "SELECT TIMESTAMP '1992-02-01 00:00 UTC' AT TIME ZONE 'Asia/Tokyo'"
+  }
 }