ClickHouse: Add python module (#418)

* ClickHouse: Update python bindings

* Update comment

* Add clickhouse to default feature

Author: trueleo

README.md

@@ -23,6 +23,7 @@ let ctx = SessionContext::with_state(state);
 - PostgreSQL
 - MySQL
 - SQLite
+- ClickHouse
 - DuckDB
 - Flight SQL
 - ODBC
@@ -86,6 +87,59 @@ cargo run -p datafusion-table-providers --example postgres --features postgres
 
 ```
 
+### ClickHouse
+
+In order to run the Clickhouse example, you need to have a Clickhouse server running. You can use the following command to start a Clickhouse server in a Docker container the example can use:
+
+```bash
+docker run --name clickhouse \
+    -e CLICKHOUSE_DB=default \
+    -e CLICKHOUSE_USER=admin \
+    -e CLICKHOUSE_PASSWORD=secret \
+    -p 8123:8123 \
+    -p 9000:9000 \
+    -d clickhouse/clickhouse-server:24.8-alpine
+
+# 2. Wait for readiness
+echo "Waiting for ClickHouse to start..."
+until curl -s http://localhost:8123/ping | grep -q 'Ok'; do
+    sleep 2
+done
+echo
+
+# 3. Create tables and a parameterized view
+docker exec -i clickhouse clickhouse-client \
+    --user=admin --password=secret --multiquery <<EOF
+CREATE TABLE IF NOT EXISTS Reports (
+  id UInt32,
+  name String
+) ENGINE = Memory;
+
+INSERT INTO Reports VALUES (1, 'Monthly Report'), (2, 'Quarterly Report');
+
+CREATE TABLE IF NOT EXISTS Users (
+  workspace_uid String,
+  id UInt32,
+  username String
+) ENGINE = Memory;
+
+INSERT INTO Users VALUES
+  ('abc', 1, 'alice'),
+  ('abc', 2, 'bob'),
+  ('xyz', 3, 'charlie');
+
+CREATE OR REPLACE VIEW Users AS
+SELECT workspace_uid, id, username
+FROM Users
+WHERE workspace_uid = {workspace_uid:String};
+EOF
+```
+
+```bash
+# Run from repo folder
+cargo run -p datafusion-table-providers --example clickhouse --features clickhouse
+```
+
 ### MySQL
 
 In order to run the MySQL example, you need to have a MySQL server running. You can use the following command to start a MySQL server in a Docker container the example can use:

core/src/sql/db_connection_pool/dbconnection/clickhouseconn.rs

@@ -49,7 +49,7 @@ impl AsyncDbConnection<Client, ()> for Client {
         }
 
         let tables: Vec<Row> = self
-            .query("SHOW TABLES FROM ?")
+            .query("SELECT name FROM system.tables WHERE database = ?")
             .bind(schema)
             .fetch_all()
             .await
@@ -65,7 +65,7 @@ impl AsyncDbConnection<Client, ()> for Client {
             name: String,
         }
         let tables: Vec<Row> = self
-            .query("SHOW DATABASES")
+            .query("SELECT name FROM system.databases WHERE name NOT IN ('system', 'information_schema', 'INFORMATION_SCHEMA')")
             .fetch_all()
             .await
             .boxed()

python/Cargo.toml

@@ -15,16 +15,30 @@ doc = false
 
 [dependencies]
 arrow = { workspace = true }
-arrow-flight = {workspace = true, optional = true}
+arrow-flight = { workspace = true, optional = true }
 datafusion = { workspace = true, features = ["pyarrow"] }
 datafusion-ffi = { workspace = true }
 datafusion-table-providers = { workspace = true }
 pyo3 = { version = "0.24.2" }
-tokio = { version = "1.46", features = ["macros", "rt", "rt-multi-thread", "sync"] }
-duckdb = { workspace = true, optional = true}
+tokio = { version = "1.46", features = [
+    "macros",
+    "rt",
+    "rt-multi-thread",
+    "sync",
+] }
+duckdb = { workspace = true, optional = true }
 
 [features]
-default = ["duckdb", "sqlite", "mysql", "postgres", "odbc", "flight"]
+default = [
+    "duckdb",
+    "clickhouse",
+    "sqlite",
+    "mysql",
+    "postgres",
+    "odbc",
+    "flight",
+]
+clickhouse = ["datafusion-table-providers/clickhouse-federation"]
 duckdb = ["dep:duckdb", "datafusion-table-providers/duckdb-federation"]
 sqlite = ["datafusion-table-providers/sqlite-federation"]
 mysql = ["datafusion-table-providers/mysql-federation"]

python/examples/clickhouse_demo.py

@@ -0,0 +1,17 @@
+from datafusion import SessionContext
+from datafusion_table_providers import clickhouse
+
+ctx = SessionContext()
+connection_param = {
+    "url": "http://localhost:8123",
+    "database": "default",
+    "user": "user",
+    "password": "secret"
+}
+pool = clickhouse.ClickHouseTableFactory(connection_param)
+tables = pool.tables()
+
+for t in tables:
+    ctx.register_table_provider(t, pool.get_table(t))
+    print("Checking table:", t)
+    ctx.table(t).show()

python/python/datafusion_table_providers/clickhouse.py

@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Python interface for Clickhouse table provider."""
+
+from typing import Any, List
+from . import _internal
+
+
+class ClickHouseTableFactory:
+    """ClickHouse table factory."""
+
+    def __init__(self, params: dict) -> None:
+        """Create a ClickHouse table factory."""
+        self._raw = _internal.clickhouse.RawClickHouseTableFactory(params)
+
+    def tables(self) -> List[str]:
+        """Get all the table names."""
+        return self._raw.tables()
+
+    def get_table(self, table_reference: str, args=None) -> Any:
+        """Return the table provider for table named `table_reference`.
+
+        Args:
+            table_reference (str): table name
+            args: optional list of parameter tuples (name, value)
+        """
+        return self._raw.get_table(table_reference, args)

python/python/tests/test_clickhouse.py

@@ -0,0 +1,89 @@
+import subprocess
+import time
+from datafusion import SessionContext
+from datafusion_table_providers import clickhouse  # hypothetical provider
+
+def run_docker_container():
+    subprocess.run(
+        ["docker", "run", "--name", "clickhouse",
+        "-e", "CLICKHOUSE_USER=user",
+        "-e", "CLICKHOUSE_PASSWORD=secret",
+        "-p", "8123:8123",
+        "-d", "clickhouse/clickhouse-server:latest"],
+        check=True,
+    )
+    time.sleep(20)
+
+def create_schema():
+    sql = r"""
+        CREATE TABLE companies (
+            id UInt32,
+            name String,
+            founded Date,
+            revenue Decimal(18,2),
+            is_active Bool,
+            tags Array(String)
+        ) ENGINE = MergeTree()
+        ORDER BY id;
+
+        INSERT INTO companies VALUES
+            (1, 'Acme Corporation', '1999-03-12', 12500000.50, 1, ['manufacturing', 'global']),
+            (2, 'Widget Inc.',      '2005-07-01',  4500000.00, 1, ['gadgets', 'innovation']),
+            (3, 'Gizmo Corp.',       '2010-11-23',   780000.75, 0, ['hardware']),
+            (4, 'Tech Solutions',    '2018-01-15',   220000.00, 1, ['consulting','it']),
+            (5, 'Data Innovations',  '2021-05-10',    98000.99, 1, ['analytics','startup']);
+
+        CREATE VIEW companies_param_view AS
+        SELECT id, name, founded, revenue, is_active, tags
+        FROM companies
+        WHERE name = {name:String};
+    """
+    subprocess.run(
+        ["docker", "exec", "-i", "clickhouse", "clickhouse-client", "--multiquery", "--query", sql],
+        check=True
+    )
+
+def stop_container():
+    subprocess.run(["docker", "stop", "clickhouse"], check=True)
+    subprocess.run(["docker", "rm",   "clickhouse"], check=True)
+
+class TestClickHouseParameterized:
+    @classmethod
+    def setup_class(cls):
+        run_docker_container()
+        create_schema()
+        cls.ctx = SessionContext()
+        connection_param = {
+            "url": "http://localhost:8123",
+            "database": "default",
+            "user": "user",
+            "password": "secret"
+        }
+        cls.pool = clickhouse.ClickHouseTableFactory(connection_param)
+
+    @classmethod
+    def teardown_class(cls):
+        stop_container()
+
+    def test_get_tables(self):
+        tables = self.pool.tables()
+        assert "companies" in tables
+        assert "companies_param_view" in tables
+
+    def test_parameterized_view(self):
+        # Register provider so DF can access the view\
+
+        self.ctx.register_table_provider(
+            "companies_param_view",
+            self.pool.get_table("companies_param_view", [("name", "Gizmo Corp.")])
+        )
+
+        # Query using parameter
+        df = self.ctx.sql(
+            "SELECT id, name, revenue FROM companies_param_view"
+        )
+        rows = df.collect()[0]
+
+        assert len(rows["name"]) == 1
+        assert rows["name"][0].as_py() == "Gizmo Corp."
+        assert rows["revenue"][0].as_py() == 780000.75

python/src/clickhouse.rs

@@ -0,0 +1,113 @@
+use std::sync::Arc;
+
+use datafusion_table_providers::{
+    clickhouse::{Arg, ClickHouseTableFactory},
+    sql::db_connection_pool::{clickhousepool::ClickHouseConnectionPool, DbConnectionPool},
+    util::secrets::to_secret_map,
+};
+use pyo3::{
+    exceptions::PyTypeError,
+    prelude::*,
+    types::{PyDict, PyList},
+};
+
+use crate::{
+    utils::{pydict_to_hashmap, to_pyerr, wait_for_future},
+    RawTableProvider,
+};
+
+#[pyclass(module = "datafusion_table_providers._internal.clickhouse")]
+struct RawClickHouseTableFactory {
+    pool: Arc<ClickHouseConnectionPool>,
+    factory: ClickHouseTableFactory,
+}
+
+#[pymethods]
+impl RawClickHouseTableFactory {
+    #[new]
+    #[pyo3(signature = (params))]
+    pub fn new(py: Python, params: &Bound<'_, PyDict>) -> PyResult<Self> {
+        let params = to_secret_map(pydict_to_hashmap(params)?);
+        let pool =
+            Arc::new(wait_for_future(py, ClickHouseConnectionPool::new(params)).map_err(to_pyerr)?);
+
+        Ok(Self {
+            factory: ClickHouseTableFactory::new(Arc::clone(&pool)),
+            pool,
+        })
+    }
+
+    pub fn tables(&self, py: Python) -> PyResult<Vec<String>> {
+        wait_for_future(py, async {
+            let conn = self.pool.connect().await.map_err(to_pyerr)?;
+            let conn_async = conn.as_async().ok_or(to_pyerr(
+                "Unable to create connection to Postgres db".to_string(),
+            ))?;
+            let schemas = conn_async.schemas().await.map_err(to_pyerr)?;
+
+            let mut tables = Vec::default();
+            for schema in schemas {
+                let schema_tables = conn_async.tables(&schema).await.map_err(to_pyerr)?;
+                tables.extend(schema_tables);
+            }
+
+            Ok(tables)
+        })
+    }
+
+    #[pyo3(signature = (table_reference, args=None))]
+    pub fn get_table(
+        &self,
+        py: Python,
+        table_reference: &str,
+        args: Option<Py<PyAny>>,
+    ) -> PyResult<RawTableProvider> {
+        let args_vec = if let Some(args) = args {
+            let seq = args.downcast_bound::<PyList>(py).map_err(|_| {
+                PyTypeError::new_err("Argument must be list of int (signed/unsigned) or string")
+            })?;
+
+            let arr: Result<Vec<_>, _> = seq
+                .iter()
+                .map(|val| {
+                    val.extract::<(String, u64)>()
+                        .map(|x| (x.0, Arg::Unsigned(x.1)))
+                        .or_else(|_| {
+                            val.extract::<(String, i64)>()
+                                .map(|x| (x.0, Arg::Signed(x.1)))
+                        })
+                        .or_else(|_| {
+                            val.extract::<(String, String)>()
+                                .map(|x| (x.0, Arg::String(x.1)))
+                        })
+                })
+                .collect();
+
+            let arr = arr.map_err(|_| {
+                PyTypeError::new_err("Argument must be list of int (signed/unsigned) or string")
+            })?;
+
+            Some(arr)
+        } else {
+            None
+        };
+
+        let table = wait_for_future(
+            py,
+            self.factory
+                .table_provider(table_reference.into(), args_vec),
+        )
+        .map_err(to_pyerr)?;
+
+        Ok(RawTableProvider {
+            table,
+            supports_pushdown_filters: true,
+        })
+    }
+}
+
+pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<RawClickHouseTableFactory>()?;
+
+    Ok(())
+}