add array join and final in sqa

Author: joe-clickhouse

clickhouse_connect/cc_sqlalchemy/__init__.py

@@ -1,6 +1,10 @@
 from clickhouse_connect import driver_name
 from clickhouse_connect.cc_sqlalchemy.datatypes.base import schema_types
+from clickhouse_connect.cc_sqlalchemy.sql import final
+from clickhouse_connect.cc_sqlalchemy.sql.clauses import array_join, ArrayJoin
 
 # pylint: disable=invalid-name
 dialect_name = driver_name
 ischema_names = schema_types
+
+__all__ = ['dialect_name', 'ischema_names', 'array_join', 'ArrayJoin', 'final']

clickhouse_connect/cc_sqlalchemy/sql/__init__.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 from sqlalchemy import Table
+from sqlalchemy.sql.selectable import FromClause, Select
 
 from clickhouse_connect.driver.binding import quote_identifier
 
@@ -13,3 +14,46 @@ def full_table(table_name: str, schema: Optional[str] = None) -> str:
 
 def format_table(table: Table):
     return full_table(table.name, table.schema)
+
+
+def final(select_stmt: Select, table: Optional[FromClause] = None) -> Select:
+    """
+    Apply the ClickHouse FINAL modifier to a select statement.
+
+    Args:
+        select_stmt: The SQLAlchemy Select statement to modify.
+        table: Optional explicit table/alias to apply FINAL to. When omitted the
+            method will use the single FROM element present on the select. A
+            ValueError is raised if the statement has no FROMs or more than one
+            FROM element and table is not provided.
+
+    Returns:
+        A new Select that renders the FINAL modifier for the target table.
+    """
+    if not isinstance(select_stmt, Select):
+        raise TypeError("final() expects a SQLAlchemy Select instance")
+
+    target = table
+    if target is None:
+        froms = select_stmt.get_final_froms()
+        if not froms:
+            raise ValueError("final() requires a table to apply the FINAL modifier.")
+        if len(froms) > 1:
+            raise ValueError("final() is ambiguous for statements with multiple FROM clauses. Specify the table explicitly.")
+        target = froms[0]
+
+    if not isinstance(target, FromClause):
+        raise TypeError("table must be a SQLAlchemy FromClause when provided")
+
+    return select_stmt.with_hint(target, "FINAL")
+
+
+def _select_final(self: Select, table: Optional[FromClause] = None) -> Select:
+    """
+    Select.final() convenience wrapper around the module-level final() helper.
+    """
+    return final(self, table=table)
+
+
+# Monkey-patch the Select class to add the .final() convenience method
+Select.final = _select_final

clickhouse_connect/cc_sqlalchemy/sql/clauses.py

@@ -0,0 +1,93 @@
+from sqlalchemy.sql.base import Immutable
+from sqlalchemy.sql.selectable import FromClause
+
+
+# pylint: disable=protected-access,too-many-ancestors,abstract-method,unused-argument
+class ArrayJoin(Immutable, FromClause):
+    """Represents ClickHouse ARRAY JOIN clause"""
+
+    __visit_name__ = "array_join"
+    _is_from_container = True
+    named_with_column = False
+    _is_join = True
+
+    def __init__(self, left, array_column, alias=None, is_left=False):
+        """Initialize ARRAY JOIN clause
+
+        Args:
+            left: The left side (table or subquery)
+            array_column: The array column to join
+            alias: Optional alias for the joined array elements
+            is_left: If True, use LEFT ARRAY JOIN instead of ARRAY JOIN
+        """
+        super().__init__()
+        self.left = left
+        self.array_column = array_column
+        self.alias = alias
+        self.is_left = is_left
+        self._is_clone_of = None
+
+    @property
+    def selectable(self):
+        """Return the selectable for this clause"""
+        return self.left
+
+    @property
+    def _hide_froms(self):
+        """Hide the left table from the FROM clause since it's part of the ARRAY JOIN"""
+        return [self.left]
+
+    @property
+    def _from_objects(self):
+        """Return all FROM objects referenced by this construct"""
+        return self.left._from_objects
+
+    def _clone(self, **kw):
+        """Return a copy of this ArrayJoin"""
+        c = self.__class__.__new__(self.__class__)
+        c.__dict__ = self.__dict__.copy()
+        c._is_clone_of = self
+        return c
+
+    def _copy_internals(self, clone=None, **kw):
+        """Copy internal state for cloning
+
+        This ensures that when queries are cloned (e.g., for subqueries, unions, or CTEs),
+        the left FromClause and array_column references are properly deep-cloned.
+        """
+        def _default_clone(elem, **kwargs):
+            return elem
+
+        if clone is None:
+            clone = _default_clone
+
+        # Clone the left FromClause and array column to ensure proper
+        #  reference handling in complex query scenarios
+        self.left = clone(self.left, **kw)
+        self.array_column = clone(self.array_column, **kw)
+
+
+def array_join(left, array_column, alias=None, is_left=False):
+    """Create an ARRAY JOIN clause
+
+    Args:
+        left: The left side (table or subquery)
+        array_column: The array column to join
+        alias: Optional alias for the joined array elements
+        is_left: If True, use LEFT ARRAY JOIN instead of ARRAY JOIN
+
+    Returns:
+        ArrayJoin: An ArrayJoin clause element
+
+    Example:
+        from clickhouse_connect.cc_sqlalchemy.sql.clauses import array_join
+
+        # Basic ARRAY JOIN
+        query = select(table).select_from(array_join(table, table.c.tags))
+
+        # LEFT ARRAY JOIN with alias
+        query = select(table).select_from(
+            array_join(table, table.c.tags, alias='tag', is_left=True)
+        )
+    """
+    return ArrayJoin(left, array_column, alias, is_left)

clickhouse_connect/cc_sqlalchemy/sql/compiler.py

@@ -1,13 +1,14 @@
 from sqlalchemy.exc import CompileError
 from sqlalchemy.sql.compiler import SQLCompiler
 
+from clickhouse_connect.cc_sqlalchemy import ArrayJoin
 from clickhouse_connect.cc_sqlalchemy.sql import format_table
 
 
 # pylint: disable=arguments-differ
 class ChStatementCompiler(SQLCompiler):
 
-    # pylint: disable=attribute-defined-outside-init
+    # pylint: disable=attribute-defined-outside-init,unused-argument
     def visit_delete(self, delete_stmt, visiting_cte=None, **kw):
         table = delete_stmt.table
         text = f"DELETE FROM {format_table(table)}"
@@ -23,10 +24,20 @@ def visit_delete(self, delete_stmt, visiting_cte=None, **kw):
 
         return text
 
-    def visit_select(self, select_stmt, **kw):
-        return super().visit_select(select_stmt, **kw)
+    def visit_array_join(self, array_join_clause, asfrom=False, from_linter=None, **kw):
+        left = self.process(array_join_clause.left, asfrom=True, from_linter=from_linter, **kw)
+        array_col = self.process(array_join_clause.array_column, **kw)
+        join_type = "LEFT ARRAY JOIN" if array_join_clause.is_left else "ARRAY JOIN"
+        text = f"{left} {join_type} {array_col}"
+        if array_join_clause.alias:
+            text += f" AS {self.preparer.quote(array_join_clause.alias)}"
+
+        return text
 
     def visit_join(self, join, **kw):
+        if isinstance(join, ArrayJoin):
+            return self.visit_array_join(join, **kw)
+
         left = self.process(join.left, **kw)
         right = self.process(join.right, **kw)
         onclause = join.onclause
@@ -60,7 +71,7 @@ def visit_column(self, column, add_to_result_map=None, include_table=True, resul
         )
 
     # Abstract methods required by SQLCompiler
-    def delete_extra_from_clause(self, delete_stmt, from_table, extra_froms, from_hints, **kw):
+    def delete_extra_from_clause(self, update_stmt, from_table, extra_froms, from_hints, **kw):
         raise NotImplementedError("ClickHouse doesn't support DELETE with extra FROM clause")
 
     def update_from_clause(self, update_stmt, from_table, extra_froms, from_hints, **kw):

tests/integration_tests/test_sqlalchemy/test_array_join.py

@@ -0,0 +1,119 @@
+from sqlalchemy import Column, MetaData, Table, literal_column, select
+from sqlalchemy.engine.base import Engine
+from sqlalchemy.types import Integer, String
+
+from clickhouse_connect.cc_sqlalchemy.datatypes.sqltypes import Array
+from clickhouse_connect.cc_sqlalchemy.datatypes.sqltypes import String as ChString
+from clickhouse_connect.cc_sqlalchemy.ddl.tableengine import MergeTree
+from clickhouse_connect.cc_sqlalchemy.sql.clauses import array_join
+
+
+def test_array_join(test_engine: Engine, test_db: str):
+    """Test ARRAY JOIN clause"""
+    with test_engine.begin() as conn:
+        metadata = MetaData(schema=test_db)
+
+        test_table = Table(
+            "test_array_join",
+            metadata,
+            Column("id", Integer),
+            Column("name", String),
+            Column("tags", Array(ChString)),
+            MergeTree(order_by="id"),
+        )
+
+        test_table.drop(conn, checkfirst=True)
+        test_table.create(conn)
+
+        conn.execute(
+            test_table.insert(),
+            [
+                {"id": 1, "name": "Alice", "tags": ["python", "sql", "clickhouse"]},
+                {"id": 2, "name": "Bob", "tags": ["java", "sql"]},
+                {"id": 3, "name": "Joe", "tags": ["python", "javascript"]},
+                {"id": 4, "name": "Charlie", "tags": []},
+            ],
+        )
+
+        query = (
+            select(test_table.c.id, test_table.c.name, test_table.c.tags)
+            .select_from(array_join(test_table, test_table.c.tags))
+            .order_by(test_table.c.id)
+            .order_by(test_table.c.tags)
+        )
+
+        compiled = query.compile(dialect=test_engine.dialect)
+        assert "ARRAY JOIN" in str(compiled).upper()
+
+        result = conn.execute(query)
+        rows = result.fetchall()
+        assert len(rows) == 7
+        assert rows[0].id == 1
+        assert rows[0].name == "Alice"
+        assert rows[0].tags == "clickhouse"
+        # ARRAY JOIN should not contain items with empty lists
+        assert "Charlie" not in [row.name for row in rows]
+
+        test_table.drop(conn)
+
+
+def test_left_array_join_with_alias(test_engine: Engine, test_db: str):
+    """Test LEFT ARRAY JOIN with alias"""
+    with test_engine.begin() as conn:
+        metadata = MetaData(schema=test_db)
+
+        test_table = Table(
+            "test_left_array_join",
+            metadata,
+            Column("id", Integer),
+            Column("name", String),
+            Column("tags", Array(ChString)),
+            MergeTree(order_by="id"),
+        )
+
+        test_table.drop(conn, checkfirst=True)
+        test_table.create(conn)
+
+        conn.execute(
+            test_table.insert(),
+            [
+                {"id": 1, "name": "Alice", "tags": ["python", "sql", "clickhouse"]},
+                {"id": 2, "name": "Bob", "tags": ["java", "sql"]},
+                {"id": 3, "name": "Joe", "tags": ["python", "javascript"]},
+                {"id": 4, "name": "Charlie", "tags": []},
+            ],
+        )
+
+        query = (
+            select(
+                test_table.c.id,
+                test_table.c.name,
+                literal_column("tag"),  # Needed when using alias
+            )
+            .select_from(array_join(test_table, test_table.c.tags, alias="tag", is_left=True))
+            .order_by(test_table.c.id)
+            .order_by(literal_column("tag"))
+        )
+
+        compiled = query.compile(dialect=test_engine.dialect)
+        compiled_str = str(compiled).upper()
+        assert "LEFT ARRAY JOIN" in compiled_str
+        assert "AS" in compiled_str
+
+        result = conn.execute(query)
+        rows = result.fetchall()
+        assert len(rows) == 8
+
+        alice_tags = [row.tag for row in rows if row.name == "Alice"]
+        assert len(alice_tags) == 3
+        assert alice_tags == sorted(["python", "sql", "clickhouse"])
+
+        bob_tags = [row.tag for row in rows if row.name == "Bob"]
+        assert len(bob_tags) == 2
+        assert bob_tags == sorted(["java", "sql"])
+
+        charlie_rows = [row for row in rows if row.name == "Charlie"]
+        assert len(charlie_rows) == 1
+        assert charlie_rows[0].tag == ""
+
+        test_table.drop(conn)