From 3518b78534dcb005df1c6a539dea218eece7192b Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Fri, 20 Oct 2023 17:13:03 -0400
Subject: [PATCH] feat(duckdb): add support for specific timestamp scales

---
 ibis/backends/base/sqlglot/datatypes.py      | 20 ++++++++++++++++++++
 ibis/backends/duckdb/tests/test_client.py    | 19 +++++++++++++++++++
 ibis/backends/duckdb/tests/test_datatypes.py |  7 +++++--
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/ibis/backends/base/sqlglot/datatypes.py b/ibis/backends/base/sqlglot/datatypes.py
index 5665783a847e..6b20645e02dd 100644
--- a/ibis/backends/base/sqlglot/datatypes.py
+++ b/ibis/backends/base/sqlglot/datatypes.py
@@ -352,6 +352,26 @@ class DuckDBType(SqlglotType):
     default_decimal_scale = 3
     default_interval_precision = "us"
 
+    @classmethod
+    def _from_sqlglot_TIMESTAMP(cls) -> dt.Timestamp:
+        return dt.Timestamp(scale=6, nullable=cls.default_nullable)
+
+    @classmethod
+    def _from_sqlglot_TIMESTAMPTZ(cls) -> dt.Timestamp:
+        return dt.Timestamp(scale=6, timezone="UTC", nullable=cls.default_nullable)
+
+    @classmethod
+    def _from_sqlglot_TIMESTAMP_S(cls) -> dt.Timestamp:
+        return dt.Timestamp(scale=0, nullable=cls.default_nullable)
+
+    @classmethod
+    def _from_sqlglot_TIMESTAMP_MS(cls) -> dt.Timestamp:
+        return dt.Timestamp(scale=3, nullable=cls.default_nullable)
+
+    @classmethod
+    def _from_sqlglot_TIMESTAMP_NS(cls) -> dt.Timestamp:
+        return dt.Timestamp(scale=9, nullable=cls.default_nullable)
+
 
 class TrinoType(SqlglotType):
     dialect = "trino"
diff --git a/ibis/backends/duckdb/tests/test_client.py b/ibis/backends/duckdb/tests/test_client.py
index f940602df84e..da472e9750bb 100644
--- a/ibis/backends/duckdb/tests/test_client.py
+++ b/ibis/backends/duckdb/tests/test_client.py
@@ -3,9 +3,12 @@
 import duckdb
 import pytest
 import sqlalchemy as sa
+from pytest import param
 
 import ibis
+import ibis.expr.datatypes as dt
 from ibis.conftest import LINUX, SANDBOXED
+from ibis.util import gen_name
 
 
 @pytest.fixture(scope="session")
@@ -113,3 +116,19 @@ def test_attach_detach(tmpdir):
 
     with pytest.raises(sa.exc.ProgrammingError):
         con2.detach(name)
+
+
+@pytest.mark.parametrize(
+    "scale",
+    [
+        None,
+        param(0, id="seconds"),
+        param(3, id="millis"),
+        param(6, id="micros"),
+        param(9, id="nanos"),
+    ],
+)
+def test_create_table_with_timestamp_scales(con, scale):
+    schema = ibis.schema(dict(ts=dt.Timestamp(scale=scale)))
+    t = con.create_table(gen_name("duckdb_timestamp_scale"), schema=schema, temp=True)
+    assert t.schema() == schema
diff --git a/ibis/backends/duckdb/tests/test_datatypes.py b/ibis/backends/duckdb/tests/test_datatypes.py
index 5ef09b45397f..d2e3bfc97c8f 100644
--- a/ibis/backends/duckdb/tests/test_datatypes.py
+++ b/ibis/backends/duckdb/tests/test_datatypes.py
@@ -29,8 +29,8 @@
             ("SMALLINT", dt.int16),
             ("TIME", dt.time),
             ("TIME WITH TIME ZONE", dt.time),
-            ("TIMESTAMP", dt.timestamp),
-            ("TIMESTAMP WITH TIME ZONE", dt.Timestamp("UTC")),
+            ("TIMESTAMP", dt.Timestamp(scale=6)),
+            ("TIMESTAMP WITH TIME ZONE", dt.Timestamp(scale=6, timezone="UTC")),
             ("TINYINT", dt.int8),
             ("UBIGINT", dt.uint64),
             ("UINTEGER", dt.uint32),
@@ -53,6 +53,9 @@
             ("INTEGER[][]", dt.Array(dt.Array(dt.int32))),
             ("JSON", dt.json),
             ("HUGEINT", dt.Decimal(38, 0)),
+            ("TIMESTAMP_S", dt.Timestamp(scale=0)),
+            ("TIMESTAMP_MS", dt.Timestamp(scale=3)),
+            ("TIMESTAMP_NS", dt.Timestamp(scale=9)),
         ]
     ],
 )