From 7826ecdb913f5aea146c2cc4a4c2f49858ed6e84 Mon Sep 17 00:00:00 2001
From: Greg Rychlewski <greg.rychlewski@gmail.com>
Date: Tue, 2 Jul 2024 09:36:33 -0400
Subject: [PATCH] Use type modifier in time and timetz extensions (#685)

---
 lib/postgrex/extensions/time.ex   | 12 +++++++++---
 lib/postgrex/extensions/timetz.ex | 15 ++++++++++-----
 test/calendar_test.exs            | 10 ++++++++++
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/lib/postgrex/extensions/time.ex b/lib/postgrex/extensions/time.ex
index de1fb25c..3eba709f 100644
--- a/lib/postgrex/extensions/time.ex
+++ b/lib/postgrex/extensions/time.ex
@@ -3,6 +3,8 @@ defmodule Postgrex.Extensions.Time do
   import Postgrex.BinaryUtils, warn: false
   use Postgrex.BinaryExtension, send: "time_send"
 
+  @default_precision 6
+
   def encode(_) do
     quote location: :keep do
       %Time{calendar: Calendar.ISO} = time ->
@@ -16,7 +18,7 @@ defmodule Postgrex.Extensions.Time do
   def decode(_) do
     quote location: :keep do
       <<8::int32(), microsecs::int64()>> ->
-        unquote(__MODULE__).microsecond_to_elixir(microsecs)
+        unquote(__MODULE__).microsecond_to_elixir(microsecs, var!(mod))
     end
   end
 
@@ -28,12 +30,16 @@ defmodule Postgrex.Extensions.Time do
     <<8::int32(), :calendar.time_to_seconds(time) * 1_000_000 + usec::int64()>>
   end
 
-  def microsecond_to_elixir(microsec) do
+  def microsecond_to_elixir(microsec, 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
+
     sec = div(microsec, 1_000_000)
     microsec = rem(microsec, 1_000_000)
 
     sec
     |> :calendar.seconds_to_time()
-    |> Time.from_erl!({microsec, 6})
+    |> Time.from_erl!({microsec, precision})
   end
 end
diff --git a/lib/postgrex/extensions/timetz.ex b/lib/postgrex/extensions/timetz.ex
index b6d476b7..9f8b34c8 100644
--- a/lib/postgrex/extensions/timetz.ex
+++ b/lib/postgrex/extensions/timetz.ex
@@ -4,6 +4,7 @@ defmodule Postgrex.Extensions.TimeTZ do
   use Postgrex.BinaryExtension, send: "timetz_send"
 
   @day (:calendar.time_to_seconds({23, 59, 59}) + 1) * 1_000_000
+  @default_precision 6
 
   def encode(_) do
     quote location: :keep do
@@ -18,7 +19,7 @@ defmodule Postgrex.Extensions.TimeTZ do
   def decode(_) do
     quote location: :keep do
       <<12::int32(), microsecs::int64(), tz::int32()>> ->
-        unquote(__MODULE__).microsecond_to_elixir(microsecs, tz)
+        unquote(__MODULE__).microsecond_to_elixir(microsecs, var!(mod), tz)
     end
   end
 
@@ -43,18 +44,22 @@ defmodule Postgrex.Extensions.TimeTZ do
     <<12::int32(), :calendar.time_to_seconds(time) * 1_000_000 + usec::int64(), 0::int32()>>
   end
 
-  def microsecond_to_elixir(microsec, tz) do
+  def microsecond_to_elixir(microsec, precision, tz) do
     microsec
     |> adjust_microsecond(tz)
-    |> microsecond_to_elixir()
+    |> microsecond_to_elixir(precision)
   end
 
-  defp microsecond_to_elixir(microsec) do
+  defp microsecond_to_elixir(microsec, 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
+
     sec = div(microsec, 1_000_000)
     microsec = rem(microsec, 1_000_000)
 
     sec
     |> :calendar.seconds_to_time()
-    |> Time.from_erl!({microsec, 6})
+    |> Time.from_erl!({microsec, precision})
   end
 end
diff --git a/test/calendar_test.exs b/test/calendar_test.exs
index 866a2983..274c3363 100644
--- a/test/calendar_test.exs
+++ b/test/calendar_test.exs
@@ -24,6 +24,11 @@ defmodule CalendarTest do
     assert [[~T[01:02:03.123456]]] = query("SELECT time '01:02:03.123456'", [])
   end
 
+  test "decode time with precision", context do
+    assert [[~T[00:00:00.000]]] = query("SELECT time(3) '00:00:00'", [])
+    assert [[[~T[00:00:00.000]]]] = query("SELECT ARRAY[time(3) '00:00:00']", [])
+  end
+
   test "decode timetz", context do
     assert [[~T[00:00:00.000000]]] = query("SELECT time with time zone '00:00:00 UTC'", [])
     assert [[~T[01:02:03.000000]]] = query("SELECT time with time zone '01:02:03 UTC'", [])
@@ -53,6 +58,11 @@ defmodule CalendarTest do
     assert [[~T[01:02:03.123456]]] = query("SELECT time with time zone '01:02:03.123456 UTC'", [])
   end
 
+  test "decode timetz with precision", context do
+    assert [[~T[00:00:00]]] = query("SELECT time(0) with time zone '00:00:00 UTC'", [])
+    assert [[[~T[00:00:00]]]] = query("SELECT ARRAY[time(0) with time zone '00:00:00 UTC']", [])
+  end
+
   test "decode date", context do
     assert [[~D[0001-01-01]]] = query("SELECT date '0001-01-01'", [])
     assert [[~D[0001-02-03]]] = query("SELECT date '0001-02-03'", [])