From 6a00ec0614e02191ee5a2d104f1fe288ea6b0d27 Mon Sep 17 00:00:00 2001
From: Elliana May <me@mause.me>
Date: Mon, 29 Jan 2024 21:48:51 +0800
Subject: [PATCH] feat: comment support

---
 duckdb_engine/__init__.py         |  3 ++-
 duckdb_engine/_supports.py        | 14 ++++++++++++
 duckdb_engine/tests/test_basic.py | 36 +++++++++++++++++++++++++++----
 3 files changed, 48 insertions(+), 5 deletions(-)
 create mode 100644 duckdb_engine/_supports.py

diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py
index ee44ebf6..b197bcd7 100644
--- a/duckdb_engine/__init__.py
+++ b/duckdb_engine/__init__.py
@@ -34,6 +34,7 @@
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.sql.selectable import Select
 
+from ._supports import has_comment_support
 from .config import apply_config, get_core_config
 from .datatypes import ISCHEMA_NAMES, register_extension_types
 
@@ -218,7 +219,7 @@ class Dialect(PGDialect_psycopg2):
     driver = "duckdb_engine"
     _has_events = False
     supports_statement_cache = False
-    supports_comments = False
+    supports_comments = has_comment_support()
     supports_sane_rowcount = False
     supports_server_side_cursors = False
     inspector = DuckDBInspector
diff --git a/duckdb_engine/_supports.py b/duckdb_engine/_supports.py
new file mode 100644
index 00000000..86735895
--- /dev/null
+++ b/duckdb_engine/_supports.py
@@ -0,0 +1,14 @@
+import duckdb
+
+
+def has_comment_support() -> bool:
+    """
+    See https://github.com/duckdb/duckdb/pull/10372
+    """
+    try:
+        with duckdb.connect(":memory:") as con:
+            con.execute("CREATE TABLE t (i INTEGER);")
+            con.execute("COMMENT ON TABLE t IS 'test';")
+    except duckdb.ParserException:
+        return False
+    return True
diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py
index da6742d3..cbdc77d5 100644
--- a/duckdb_engine/tests/test_basic.py
+++ b/duckdb_engine/tests/test_basic.py
@@ -35,7 +35,8 @@
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session, relationship, sessionmaker
 
-from .. import DBAPI, Dialect, supports_attach
+from .. import Dialect, supports_attach
+from .._supports import has_comment_support
 
 try:
     # sqlalchemy 2
@@ -394,10 +395,37 @@ def test_binary(session: Session) -> None:
     assert b.text == "Hello World!"
 
 
-def test_comment_support() -> None:
+@mark.skipif(
+    not has_comment_support(), reason="comments not supported by duckdb version"
+)
+def test_comment_support(engine: Engine) -> None:
     "comments not yet supported by duckdb"
-    with raises(DBAPI.ParserException, match="syntax error"):
-        duckdb.default_connection.execute('comment on sqlite_master is "hello world";')
+
+    class Notice(Base):
+        __tablename__ = "tb_notice"
+        __table_args__ = {"comment": "Notice table"}
+
+        seqno = Column(
+            Integer,
+            Sequence("seqno_sequence"),
+            primary_key=True,
+            comment="Integer representing the sequence number",
+        )
+        title = Column(
+            String(200),
+            nullable=False,
+            comment="Title of the notice, represented as a string",
+        )
+
+    Base.metadata.create_all(bind=engine)
+
+    inspector = inspect(engine)
+
+    assert inspector.get_table_comment("tb_notice")["text"] == "Notice table"
+
+    columns = inspector.get_columns("tb_notice", None)
+    assert columns[0]["comment"] == "Integer representing the sequence number"
+    assert columns[1]["comment"] == "Title of the notice, represented as a string"
 
 
 def test_rowcount() -> None: