From 3bd3770a412ae352a4d3cea775b8841aa0d419ad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Edgar=20Ram=C3=ADrez=20Mondrag=C3=B3n?=
 <16805946+edgarrmondragon@users.noreply.github.com>
Date: Thu, 18 Jul 2024 21:30:26 -0600
Subject: [PATCH] test: Add (DuckDB) SQL sink tests (#2548)

---
 samples/sample_duckdb/__init__.py  |  7 ++++
 samples/sample_duckdb/connector.py | 24 +++++++++++
 tests/core/sinks/test_sql_sink.py  | 64 ++++++++++++++++++++++++++++++
 tests/core/test_connector_sql.py   | 20 +---------
 4 files changed, 96 insertions(+), 19 deletions(-)
 create mode 100644 samples/sample_duckdb/__init__.py
 create mode 100644 samples/sample_duckdb/connector.py
 create mode 100644 tests/core/sinks/test_sql_sink.py

diff --git a/samples/sample_duckdb/__init__.py b/samples/sample_duckdb/__init__.py
new file mode 100644
index 000000000..bbfa07726
--- /dev/null
+++ b/samples/sample_duckdb/__init__.py
@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from .connector import DuckDBConnector
+
+__all__ = [
+    "DuckDBConnector",
+]
diff --git a/samples/sample_duckdb/connector.py b/samples/sample_duckdb/connector.py
new file mode 100644
index 000000000..e8c6c251f
--- /dev/null
+++ b/samples/sample_duckdb/connector.py
@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import sqlalchemy as sa
+
+from singer_sdk.connectors import SQLConnector
+
+
+class DuckDBConnector(SQLConnector):
+    allow_column_alter = True
+
+    @staticmethod
+    def get_column_alter_ddl(
+        table_name: str,
+        column_name: str,
+        column_type: sa.types.TypeEngine,
+    ) -> sa.DDL:
+        return sa.DDL(
+            "ALTER TABLE %(table_name)s ALTER COLUMN %(column_name)s TYPE %(column_type)s",  # noqa: E501
+            {
+                "table_name": table_name,
+                "column_name": column_name,
+                "column_type": column_type,
+            },
+        )
diff --git a/tests/core/sinks/test_sql_sink.py b/tests/core/sinks/test_sql_sink.py
new file mode 100644
index 000000000..c20be40f6
--- /dev/null
+++ b/tests/core/sinks/test_sql_sink.py
@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import typing as t
+from textwrap import dedent
+
+import pytest
+
+from samples.sample_duckdb import DuckDBConnector
+from singer_sdk.sinks.sql import SQLSink
+from singer_sdk.target_base import SQLTarget
+
+
+class DuckDBSink(SQLSink):
+    connector_class = DuckDBConnector
+
+
+class DuckDBTarget(SQLTarget):
+    """DuckDB target class."""
+
+    name = "sql-target-mock"
+    config_jsonschema: t.ClassVar[dict] = {"type": "object", "properties": {}}
+    default_sink_class = DuckDBSink
+
+
+class TestDuckDBSink:
+    @pytest.fixture
+    def target(self) -> DuckDBTarget:
+        return DuckDBTarget(config={"sqlalchemy_url": "duckdb:///"})
+
+    @pytest.fixture
+    def schema(self) -> dict:
+        return {
+            "properties": {
+                "id": {
+                    "type": ["string", "null"],
+                },
+                "col_ts": {
+                    "format": "date-time",
+                    "type": ["string", "null"],
+                },
+                "table": {
+                    "type": ["string", "null"],
+                },
+            },
+        }
+
+    @pytest.fixture
+    def sink(self, target: DuckDBTarget, schema: dict) -> DuckDBSink:
+        return DuckDBSink(
+            target,
+            stream_name="foo",
+            schema=schema,
+            key_properties=["id"],
+        )
+
+    def test_generate_insert_statement(self, sink: DuckDBSink, schema: dict):
+        """Test that the insert statement is generated correctly."""
+        expected = dedent(
+            """\
+            INSERT INTO foo
+            (id, col_ts, "table")
+            VALUES (:id, :col_ts, :table)"""
+        )
+        assert sink.generate_insert_statement("foo", schema=schema) == expected
diff --git a/tests/core/test_connector_sql.py b/tests/core/test_connector_sql.py
index ff3ac5735..10ee0c0f4 100644
--- a/tests/core/test_connector_sql.py
+++ b/tests/core/test_connector_sql.py
@@ -8,6 +8,7 @@
 import sqlalchemy as sa
 from sqlalchemy.dialects import registry, sqlite
 
+from samples.sample_duckdb import DuckDBConnector
 from singer_sdk.connectors import SQLConnector
 from singer_sdk.exceptions import ConfigValidationError
 
@@ -289,25 +290,6 @@ def test_engine_json_serialization(self, connector: SQLConnector):
             ]
 
 
-class DuckDBConnector(SQLConnector):
-    allow_column_alter = True
-
-    @staticmethod
-    def get_column_alter_ddl(
-        table_name: str,
-        column_name: str,
-        column_type: sa.types.TypeEngine,
-    ) -> sa.DDL:
-        return sa.DDL(
-            "ALTER TABLE %(table_name)s ALTER COLUMN %(column_name)s TYPE %(column_type)s",  # noqa: E501
-            {
-                "table_name": table_name,
-                "column_name": column_name,
-                "column_type": column_type,
-            },
-        )
-
-
 class TestDuckDBConnector:
     @pytest.fixture
     def connector(self):