Add precision to timestamptz extension (#691)

lib/postgrex/extensions/timestamptz.ex

@@ -9,6 +9,7 @@ defmodule Postgrex.Extensions.TimestampTZ do
 
   @plus_infinity 9_223_372_036_854_775_807
   @minus_infinity -9_223_372_036_854_775_808
+  @default_precision 6
 
   def init(opts), do: Keyword.get(opts, :allow_infinite_timestamps, false)
 
@@ -26,7 +27,7 @@ defmodule Postgrex.Extensions.TimestampTZ do
   def decode(infinity?) do
     quote location: :keep do
       <<8::int32(), microsecs::int64()>> ->
-        unquote(__MODULE__).microsecond_to_elixir(microsecs, unquote(infinity?))
+        unquote(__MODULE__).microsecond_to_elixir(microsecs, var!(mod), unquote(infinity?))
     end
   end
 
@@ -41,16 +42,35 @@ defmodule Postgrex.Extensions.TimestampTZ do
     raise ArgumentError, "#{inspect(datetime)} is not in UTC"
   end
 
-  def microsecond_to_elixir(@plus_infinity, infinity?) do
+  def microsecond_to_elixir(@plus_infinity, _precision, infinity?) do
     if infinity?, do: :inf, else: raise_infinity("infinity")
   end
 
-  def microsecond_to_elixir(@minus_infinity, infinity?) do
+  def microsecond_to_elixir(@minus_infinity, _precision, infinity?) do
     if infinity?, do: :"-inf", else: raise_infinity("-infinity")
   end
 
-  def microsecond_to_elixir(microsecs, _infinity) do
-    DateTime.from_unix!(microsecs + @us_epoch, :microsecond)
+  def microsecond_to_elixir(microsecs, precision, _infinity) do
+    split(microsecs, precision)
+  end
+
+  defp split(microsecs, precision) when microsecs < 0 and rem(microsecs, 1_000_000) != 0 do
+    secs = div(microsecs, 1_000_000) - 1
+    microsecs = 1_000_000 + rem(microsecs, 1_000_000)
+    split(secs, microsecs, precision)
+  end
+
+  defp split(microsecs, precision) do
+    secs = div(microsecs, 1_000_000)
+    microsecs = rem(microsecs, 1_000_000)
+    split(secs, microsecs, precision)
+  end
+
+  defp split(secs, microsecs, precision) do
+    # use the default precision if the precision modifier from postgres is -1 (this means no precision specified)
+    # or if the precision is missing because we are in a super type which does not give us the sub-type's modifier
+    precision = if precision in [-1, nil], do: @default_precision, else: precision
+    DateTime.from_gregorian_seconds(secs + @gs_epoch, {microsecs, precision})
   end
 
   defp raise_infinity(type) do

test/calendar_test.exs

@@ -104,6 +104,14 @@ defmodule CalendarTest do
              query("SELECT timestamp '4713-01-01BC 00:00:00.123456'", [])
   end
 
+  test "decode timestamp with precision", context do
+    assert [[~N[2001-01-01 00:00:00.00]]] =
+             query("SELECT timestamp(2) '2001-01-01 00:00:00'", [])
+
+    assert [[[~N[2001-01-01 00:00:00.00]]]] =
+             query("SELECT ARRAY[timestamp(2) '2001-01-01 00:00:00']", [])
+  end
+
   test "decode timestamptz", context do
     assert [
              [
@@ -222,6 +230,42 @@ defmodule CalendarTest do
            ] = query("SELECT timestamp with time zone '1980-01-01 01:00:00.123456'", [])
   end
 
+  test "decode timestamptz with precision", context do
+    assert [
+             [
+               %DateTime{
+                 year: 2001,
+                 month: 1,
+                 day: 1,
+                 hour: 0,
+                 minute: 0,
+                 second: 0,
+                 microsecond: {0, 1},
+                 time_zone: "Etc/UTC",
+                 utc_offset: 0
+               }
+             ]
+           ] = query("SELECT timestamp(1) with time zone '2001-01-01 00:00:00 UTC'", [])
+
+    assert [
+             [
+               [
+                 %DateTime{
+                   year: 2001,
+                   month: 1,
+                   day: 1,
+                   hour: 0,
+                   minute: 0,
+                   second: 0,
+                   microsecond: {0, 1},
+                   time_zone: "Etc/UTC",
+                   utc_offset: 0
+                 }
+               ]
+             ]
+           ] = query("SELECT ARRAY[timestamp(1) with time zone '2001-01-01 00:00:00 UTC']", [])
+  end
+
   test "decode negative timestampz", context do
     assert :ok = query("SET SESSION TIME ZONE UTC", [])
 

test/query_test.exs

@@ -1769,24 +1769,4 @@ defmodule QueryTest do
     {:ok, pid} = P.start_link(database: "postgrex_test", search_path: ["public", "test_schema"])
     %{rows: [[1, "foo"]]} = P.query!(pid, "SELECT * from test_table", [])
   end
-
-  test "timestamp precision", context do
-    :ok =
-      query(
-        """
-        INSERT INTO timestamps (micro, milli, sec, sec_arr)
-        VALUES ('2000-01-01', '2000-01-01', '2000-01-01', '{2000-01-01, 2000-01-02}'),
-        ('3000-01-01', '3000-01-01', '3000-01-01', '{3000-01-01, 3000-01-02}')
-        """,
-        []
-      )
-
-    assert [row1, row2] = query("SELECT * FROM timestamps", [])
-
-    assert [6, 3, 0, [0, 0]] = precision(row1)
-    assert [6, 3, 0, [0, 0]] = precision(row2)
-  end
-
-  defp precision([_ | _] = dts), do: Enum.map(dts, &precision(&1))
-  defp precision(%NaiveDateTime{microsecond: {_, p}}), do: p
 end