implement at the cpp level

Author: NickCrews

_duckdb-stubs/__init__.pyi

@@ -13,7 +13,7 @@ if typing.TYPE_CHECKING:
     import pandas
     import pyarrow.lib
     from collections.abc import Callable, Iterable, Sequence, Mapping
-    from duckdb import sqltypes, func
+    from duckdb import sqltypes, func, template
     from builtins import list as lst  # needed to avoid mypy error on DuckDBPyRelation.list method shadowing
 
     # the field_ids argument to to_parquet and write_parquet has a recursive structure
@@ -241,8 +241,12 @@ class DuckDBPyConnection:
     def dtype(self, type_str: str) -> sqltypes.DuckDBPyType: ...
     def duplicate(self) -> DuckDBPyConnection: ...
     def enum_type(self, name: str, type: sqltypes.DuckDBPyType, values: lst[typing.Any]) -> sqltypes.DuckDBPyType: ...
-    def execute(self, query: Statement | str, parameters: object = None) -> DuckDBPyConnection: ...
-    def executemany(self, query: Statement | str, parameters: object = None) -> DuckDBPyConnection: ...
+    def execute(
+        self, query: Statement | str | template.SqlTemplate | template.CompiledSql, parameters: object = None
+    ) -> DuckDBPyConnection: ...
+    def executemany(
+        self, query: Statement | str | template.SqlTemplate | template.CompiledSql, parameters: object = None
+    ) -> DuckDBPyConnection: ...
     def extract_statements(self, query: str) -> lst[Statement]: ...
     def fetch_arrow_table(self, rows_per_batch: typing.SupportsInt = 1000000) -> pyarrow.lib.Table:
         """Deprecated: use to_arrow_table() instead."""
@@ -331,7 +335,9 @@ class DuckDBPyConnection:
         union_by_name: bool = False,
         compression: str | None = None,
     ) -> DuckDBPyRelation: ...
-    def from_query(self, query: str, *, alias: str = "", params: object = None) -> DuckDBPyRelation: ...
+    def from_query(
+        self, query: str | template.SqlTemplate | template.CompiledSql, *, alias: str = "", params: object = None
+    ) -> DuckDBPyRelation: ...
     def get_table_names(self, query: str, *, qualified: bool = False) -> set[str]: ...
     def install_extension(
         self,
@@ -360,7 +366,9 @@ class DuckDBPyConnection:
     def pl(
         self, rows_per_batch: typing.SupportsInt = 1000000, *, lazy: bool = False
     ) -> polars.DataFrame | polars.LazyFrame: ...
-    def query(self, query: str, *, alias: str = "", params: object = None) -> DuckDBPyRelation: ...
+    def query(
+        self, query: str | template.SqlTemplate | template.CompiledSql, *, alias: str = "", params: object = None
+    ) -> DuckDBPyRelation: ...
     def query_progress(self) -> float: ...
     def read_csv(
         self,
@@ -462,7 +470,13 @@ class DuckDBPyConnection:
     def row_type(
         self, fields: dict[str, sqltypes.DuckDBPyType] | lst[sqltypes.DuckDBPyType]
     ) -> sqltypes.DuckDBPyType: ...
-    def sql(self, query: Statement | str, *, alias: str = "", params: object = None) -> DuckDBPyRelation: ...
+    def sql(
+        self,
+        query: Statement | str | template.SqlTemplate | template.CompiledSql,
+        *,
+        alias: str = "",
+        params: object = None,
+    ) -> DuckDBPyRelation: ...
     def sqltype(self, type_str: str) -> sqltypes.DuckDBPyType: ...
     def string_type(self, collation: str = "") -> sqltypes.DuckDBPyType: ...
     def struct_type(
@@ -1160,7 +1174,7 @@ def enum_type(
     connection: DuckDBPyConnection | None = None,
 ) -> sqltypes.DuckDBPyType: ...
 def execute(
-    query: Statement | str,
+    query: Statement | str | template.SqlTemplate | template.CompiledSql,
     parameters: object = None,
     *,
     connection: DuckDBPyConnection | None = None,
@@ -1282,7 +1296,7 @@ def from_parquet(
     connection: DuckDBPyConnection | None = None,
 ) -> DuckDBPyRelation: ...
 def from_query(
-    query: Statement | str,
+    query: Statement | str | template.SqlTemplate | template.CompiledSql,
     *,
     alias: str = "",
     params: object = None,
@@ -1350,7 +1364,7 @@ def project(
     df: pandas.DataFrame, *args: _ExpressionLike, groups: str = "", connection: DuckDBPyConnection | None = None
 ) -> DuckDBPyRelation: ...
 def query(
-    query: Statement | str,
+    query: Statement | str | template.SqlTemplate | template.CompiledSql,
     *,
     alias: str = "",
     params: object = None,
@@ -1474,7 +1488,7 @@ def row_type(
 def rowcount(*, connection: DuckDBPyConnection | None = None) -> int: ...
 def set_default_connection(connection: DuckDBPyConnection) -> None: ...
 def sql(
-    query: Statement | str,
+    query: Statement | str | template.SqlTemplate | template.CompiledSql,
     *,
     alias: str = "",
     params: object = None,

src/duckdb_py/include/duckdb_python/pyconnection/pyconnection.hpp

@@ -363,6 +363,7 @@ struct DuckDBPyConnection : public enable_shared_from_this<DuckDBPyConnection> {
 	                               FunctionNullHandling null_handling, PythonExceptionHandling exception_handling,
 	                               bool side_effects);
 	void RegisterArrowObject(const py::object &arrow_object, const string &name);
+	pair<string, py::object> ExtractCompiledSqlAndParams(const py::object &query, py::object params);
 	vector<unique_ptr<SQLStatement>> GetStatements(const py::object &query);
 
 	static PythonEnvironmentType environment;

src/duckdb_py/pyconnection.cpp

@@ -708,11 +708,74 @@ shared_ptr<DuckDBPyConnection> DuckDBPyConnection::ExecuteFromString(const strin
 	return Execute(py::str(query));
 }
 
+pair<string, py::object> DuckDBPyConnection::ExtractCompiledSqlAndParams(const py::object &query, py::object params) {
+	py::object compiled = query;
+	if (!py::hasattr(compiled, "sql") || !py::hasattr(compiled, "params")) {
+		if (!py::hasattr(query, "compile")) {
+			return {string(), params};
+		}
+		compiled = query.attr("compile")();
+	}
+
+	if (!py::hasattr(compiled, "sql") || !py::hasattr(compiled, "params")) {
+		return {string(), params};
+	}
+
+	auto compiled_sql = py::cast<string>(compiled.attr("sql"));
+	auto compiled_params_obj = compiled.attr("params");
+	if (!py::is_dict_like(compiled_params_obj)) {
+		throw InvalidInputException("Compiled SQL parameters must be a dictionary");
+	}
+
+	auto compiled_params = py::cast<py::dict>(compiled_params_obj);
+	if (compiled_params.empty()) {
+		return {compiled_sql, params};
+	}
+
+	if (params.is_none()) {
+		return {compiled_sql, compiled_params};
+	}
+
+	if (py::is_dict_like(params)) {
+		auto merged_params = py::dict();
+		for (auto &item : compiled_params) {
+			merged_params[item.first] = item.second;
+		}
+		auto provided_params = py::cast<py::dict>(params);
+		for (auto &item : provided_params) {
+			if (merged_params.contains(item.first)) {
+				throw py::value_error("Cannot merge compiled SQL parameters with duplicate parameter names");
+			}
+			merged_params[item.first] = item.second;
+		}
+		return {compiled_sql, merged_params};
+	}
+
+	if (py::is_list_like(params)) {
+		if (py::len(params) == 0) {
+			return {compiled_sql, compiled_params};
+		}
+		throw py::value_error("Cannot merge compiled SQL named parameters with positional parameters");
+	}
+
+	throw InvalidInputException("Prepared parameters can only be passed as a list or a dictionary");
+}
+
 shared_ptr<DuckDBPyConnection> DuckDBPyConnection::Execute(const py::object &query, py::object params) {
 	py::gil_scoped_acquire gil;
 	con.SetResult(nullptr);
 
-	auto statements = GetStatements(query);
+	auto normalized_query = ExtractCompiledSqlAndParams(query, params);
+	auto &compiled_sql = normalized_query.first;
+	auto &merged_params = normalized_query.second;
+	vector<unique_ptr<SQLStatement>> statements;
+	if (!compiled_sql.empty()) {
+		statements = GetStatements(py::str(compiled_sql));
+		params = merged_params;
+	} else {
+		statements = GetStatements(query);
+	}
+
 	if (statements.empty()) {
 		// TODO: should we throw?
 		return nullptr;
@@ -1603,7 +1666,17 @@ unique_ptr<DuckDBPyRelation> DuckDBPyConnection::RunQuery(const py::object &quer
 		alias = "unnamed_relation_" + StringUtil::GenerateRandomName(16);
 	}
 
-	auto statements = GetStatements(query);
+	auto normalized_query = ExtractCompiledSqlAndParams(query, params);
+	auto &compiled_sql = normalized_query.first;
+	auto &merged_params = normalized_query.second;
+	vector<unique_ptr<SQLStatement>> statements;
+	if (!compiled_sql.empty()) {
+		statements = GetStatements(py::str(compiled_sql));
+		params = merged_params;
+	} else {
+		statements = GetStatements(query);
+	}
+
 	if (statements.empty()) {
 		// TODO: should we throw?
 		return nullptr;
@@ -1616,7 +1689,7 @@ unique_ptr<DuckDBPyRelation> DuckDBPyConnection::RunQuery(const py::object &quer
 
 	// Attempt to create a Relation for lazy execution if possible
 	shared_ptr<Relation> relation;
-	bool has_params = !py::none().is(params) && py::len(params) > 0;
+	bool has_params = !params.is_none() && py::len(params) > 0;
 	if (!has_params) {
 		// No params (or empty params) — use lazy QueryRelation path
 		{

src/duckdb_py/pyexpression/initialize.cpp

@@ -314,6 +314,7 @@ void DuckDBPyExpression::Initialize(py::module_ &m) {
 		Print the stringified version of the expression.
 	)";
 	expression.def("show", &DuckDBPyExpression::Print, docs);
+	expression.def("__duckdb_template__", &DuckDBPyExpression::ToString);
 
 	docs = R"(
 		Set the order by modifier to ASCENDING.

src/duckdb_py/pyrelation/initialize.cpp

@@ -342,6 +342,7 @@ void DuckDBPyRelation::Initialize(py::handle &m) {
 	    .def("show", &DuckDBPyRelation::Print, "Display a summary of the data", py::kw_only(),
 	         py::arg("max_width") = py::none(), py::arg("max_rows") = py::none(), py::arg("max_col_width") = py::none(),
 	         py::arg("null_value") = py::none(), py::arg("render_mode") = py::none())
+	    .def("__duckdb_template__", &DuckDBPyRelation::ToSQL)
 	    .def("__str__", &DuckDBPyRelation::ToString)
 	    .def("__repr__", &DuckDBPyRelation::ToString);
 

src/duckdb_py/typing/pytype.cpp

@@ -329,6 +329,7 @@ void DuckDBPyType::Initialize(py::handle &m) {
 	auto type_module = py::class_<DuckDBPyType, shared_ptr<DuckDBPyType>>(m, "DuckDBPyType", py::module_local());
 
 	type_module.def("__repr__", &DuckDBPyType::ToString, "Stringified representation of the type object");
+	type_module.def("__duckdb_template__", &DuckDBPyType::ToString);
 	type_module.def("__eq__", &DuckDBPyType::Equals, "Compare two types for equality", py::arg("other"),
 	                py::is_operator());
 	type_module.def("__eq__", &DuckDBPyType::EqualsString, "Compare two types for equality", py::arg("other"),

tests/fast/test_template_e2e.py

@@ -41,8 +41,8 @@ def test_module_level_sql_apis_accept_sql_template() -> None:
 
 def test_module_level_execute_accepts_sql_template() -> None:
     conn = duckdb.connect()
-    query = template("SELECT ", "hello")
-    assert duckdb.execute(query, connection=conn).fetchone() == ("hello",)
+    query = template("SELECT ", 5)
+    assert duckdb.execute(query, connection=conn).fetchone() == (5,)
 
 
 def test_connection_sql_accepts_alias_kwarg_with_template() -> None:
@@ -61,11 +61,19 @@ def test_connection_sql_template_can_merge_additional_params() -> None:
 
 def test_connection_sql_template_param_name_conflict_with_additional_params_raises() -> None:
     conn = duckdb.connect()
-    query = template("SELECT ", param(10, "num"), " + $num")
+    query = template("SELECT ", param(10, "num", exact=True), " + $num")
     with pytest.raises((duckdb.InvalidInputException, ValueError)):
         conn.sql(query, params={"num": 5}).fetchall()
 
 
+def test_cant_merge_with_positional_params() -> None:
+    conn = duckdb.connect()
+    # It doesn't even have a name, but still should error
+    query = template("SELECT ", 10, " + ?")
+    with pytest.raises(ValueError, match="Cannot merge compiled SQL named parameters with positional parameters"):
+        conn.sql(query, params=[5]).fetchall()
+
+
 def test_sql_apis_accept_compiled_sql() -> None:
     conn = duckdb.connect()
     compiled = template("SELECT i FROM range(5) t(i) WHERE i >= ", 3, " ORDER BY i").compile()
@@ -83,16 +91,6 @@ def test_relation_interpolation_works_end_to_end() -> None:
     assert conn.sql(query).fetchall() == [(0,), (2,), (4,)]
 
 
-def test_interpolated_strings_are_parameterized_by_default() -> None:
-    conn = duckdb.connect()
-    conn.execute("CREATE TABLE names(name VARCHAR)")
-    conn.execute("INSERT INTO names VALUES ('alice'), ('bob')")
-
-    untrusted = "alice' OR 1=1 --"
-    query = template("SELECT count(*) FROM names WHERE name = ", untrusted)
-    assert conn.sql(query).fetchone() == (0,)
-
-
 def test_builtin_duckdbpytype_object_interpolates_in_template() -> None:
     conn = duckdb.connect()
     integer_type = duckdb.sqltype("INTEGER")