Merge pull request #1174 from Mause/upsert-conflicts

fix: reexport postgresql insert function

duckdb_engine/__init__.py

@@ -18,7 +18,7 @@
 import sqlalchemy
 from sqlalchemy import pool, select, sql, text, util
 from sqlalchemy import types as sqltypes
-from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.dialects.postgresql import UUID, insert
 from sqlalchemy.dialects.postgresql.base import (
     PGDialect,
     PGIdentifierPreparer,
@@ -53,6 +53,16 @@
 register_extension_types()
 
 
+__all__ = [
+    "Dialect",
+    "ConnectionWrapper",
+    "CursorWrapper",
+    "DBAPI",
+    "DuckDBEngineWarning",
+    "insert",  # reexport of sqlalchemy.dialects.postgresql.insert
+]
+
+
 class DBAPI:
     paramstyle = "numeric_dollar" if sqlalchemy_version >= "2.0.0" else "qmark"
     apilevel = duckdb.apilevel

duckdb_engine/tests/test_basic.py

@@ -36,7 +36,7 @@
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session, relationship, sessionmaker
 
-from .. import Dialect, supports_attach, supports_user_agent
+from .. import Dialect, insert, supports_attach, supports_user_agent
 from .._supports import has_comment_support
 
 try:
@@ -653,3 +653,29 @@ def test_reflection(engine: Engine) -> None:
     with engine.connect() as conn:
         conn.execute(text("CREATE TABLE tbl(col1 INTEGER)"))
     metadata.reflect(engine)
+
+
+def test_upsert(session: Session) -> None:
+    class User(Base):
+        __tablename__ = "users"
+
+        id = Column(Integer(), Sequence("id_seq"), primary_key=True)
+        name = Column(String, unique=True)
+        fullname = Column(String)
+
+    Base.metadata.create_all(session.bind)
+    stmt = insert(User).values(
+        [
+            {"name": "spongebob", "fullname": "Spongebob Squarepants"},
+            {"name": "sandy", "fullname": "Sandy Cheeks"},
+            {"name": "patrick", "fullname": "Patrick Star"},
+            {"name": "squidward", "fullname": "Squidward Tentacles"},
+            {"name": "ehkrabs", "fullname": "Eugene H. Krabs"},
+        ]
+    )
+    stmt = stmt.on_conflict_do_update(
+        index_elements=[User.name], set_=dict(fullname=stmt.excluded.fullname)
+    )
+    session.execute(stmt)
+
+    assert session.query(User).count() == 5