Skip to content

Commit f1141c6

Browse files
NickCrewsNickCrews
committed
implement at the cpp level
1 parent 53a3ba4 commit f1141c6

7 files changed

Lines changed: 115 additions & 26 deletions

File tree

‎_duckdb-stubs/__init__.pyi‎

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ if typing.TYPE_CHECKING:
4343
ArrowUDF,
4444
)
4545
from ._enums import ExplainTypeLiteral, RenderModeLiteral
46-
from duckdb import sqltypes, func
46+
from duckdb import sqltypes, func, template
4747

4848
__all__: lst[str] = [
4949
"BinderException",
@@ -243,8 +243,12 @@ class DuckDBPyConnection:
243243
def dtype(self, type_str: StrIntoPyType) -> sqltypes.DuckDBPyType: ...
244244
def duplicate(self) -> DuckDBPyConnection: ...
245245
def enum_type(self, name: str, type: sqltypes.DuckDBPyType, values: lst[typing.Any]) -> sqltypes.DuckDBPyType: ...
246-
def execute(self, query: Statement | str, parameters: object = None) -> DuckDBPyConnection: ...
247-
def executemany(self, query: Statement | str, parameters: object = None) -> DuckDBPyConnection: ...
246+
def execute(
247+
self, query: Statement | str | template.SqlTemplate | template.CompiledSql, parameters: object = None
248+
) -> DuckDBPyConnection: ...
249+
def executemany(
250+
self, query: Statement | str | template.SqlTemplate | template.CompiledSql, parameters: object = None
251+
) -> DuckDBPyConnection: ...
248252
def extract_statements(self, query: str) -> lst[Statement]: ...
249253
def fetch_arrow_table(self, rows_per_batch: typing.SupportsInt = 1000000) -> pyarrow.lib.Table:
250254
"""Deprecated: use to_arrow_table() instead."""
@@ -326,7 +330,9 @@ class DuckDBPyConnection:
326330
union_by_name: bool = False,
327331
compression: ParquetCompression | None = None,
328332
) -> DuckDBPyRelation: ...
329-
def from_query(self, query: str, *, alias: str = "", params: object = None) -> DuckDBPyRelation: ...
333+
def from_query(
334+
self, query: str | template.SqlTemplate | template.CompiledSql, *, alias: str = "", params: object = None
335+
) -> DuckDBPyRelation: ...
330336
def get_table_names(self, query: str, *, qualified: bool = False) -> set[str]: ...
331337
def install_extension(
332338
self,
@@ -355,7 +361,9 @@ class DuckDBPyConnection:
355361
def pl(
356362
self, rows_per_batch: typing.SupportsInt = 1000000, *, lazy: bool = False
357363
) -> polars.DataFrame | polars.LazyFrame: ...
358-
def query(self, query: str, *, alias: str = "", params: object = None) -> DuckDBPyRelation: ...
364+
def query(
365+
self, query: str | template.SqlTemplate | template.CompiledSql, *, alias: str = "", params: object = None
366+
) -> DuckDBPyRelation: ...
359367
def query_progress(self) -> float: ...
360368
def read_csv(
361369
self,
@@ -448,7 +456,13 @@ class DuckDBPyConnection:
448456
def remove_function(self, name: str) -> DuckDBPyConnection: ...
449457
def rollback(self) -> DuckDBPyConnection: ...
450458
def row_type(self, fields: IntoFields) -> sqltypes.DuckDBPyType: ...
451-
def sql(self, query: Statement | str, *, alias: str = "", params: object = None) -> DuckDBPyRelation: ...
459+
def sql(
460+
self,
461+
query: Statement | str | template.SqlTemplate | template.CompiledSql,
462+
*,
463+
alias: str = "",
464+
params: object = None,
465+
) -> DuckDBPyRelation: ...
452466
def sqltype(self, type_str: str) -> sqltypes.DuckDBPyType: ...
453467
def string_type(self, collation: str = "") -> sqltypes.DuckDBPyType: ...
454468
def struct_type(self, fields: IntoFields) -> sqltypes.DuckDBPyType: ...
@@ -952,7 +966,7 @@ def enum_type(
952966
connection: DuckDBPyConnection | None = None,
953967
) -> sqltypes.DuckDBPyType: ...
954968
def execute(
955-
query: Statement | str,
969+
query: Statement | str | template.SqlTemplate | template.CompiledSql,
956970
parameters: object = None,
957971
*,
958972
connection: DuckDBPyConnection | None = None,
@@ -1067,7 +1081,7 @@ def from_parquet(
10671081
connection: DuckDBPyConnection | None = None,
10681082
) -> DuckDBPyRelation: ...
10691083
def from_query(
1070-
query: Statement | str,
1084+
query: Statement | str | template.SqlTemplate | template.CompiledSql,
10711085
*,
10721086
alias: str = "",
10731087
params: object = None,
@@ -1132,7 +1146,7 @@ def project(
11321146
df: pandas.DataFrame, *args: IntoExpr, groups: str = "", connection: DuckDBPyConnection | None = None
11331147
) -> DuckDBPyRelation: ...
11341148
def query(
1135-
query: Statement | str,
1149+
query: Statement | str | template.SqlTemplate | template.CompiledSql,
11361150
*,
11371151
alias: str = "",
11381152
params: object = None,
@@ -1245,7 +1259,7 @@ def row_type(fields: IntoFields, *, connection: DuckDBPyConnection | None = None
12451259
def rowcount(*, connection: DuckDBPyConnection | None = None) -> int: ...
12461260
def set_default_connection(connection: DuckDBPyConnection) -> None: ...
12471261
def sql(
1248-
query: Statement | str,
1262+
query: Statement | str | template.SqlTemplate | template.CompiledSql,
12491263
*,
12501264
alias: str = "",
12511265
params: object = None,

‎src/include/duckdb_python/pyconnection/pyconnection.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,7 @@ struct DuckDBPyConnection : public std::enable_shared_from_this<DuckDBPyConnecti
376376
ScalarFunction CreateScalarUDF(const string &name, const nb::callable &udf, const nb::object &parameters,
377377
const nb::object &return_type, bool vectorized, FunctionNullHandling null_handling,
378378
PythonExceptionHandling exception_handling, bool side_effects);
379+
pair<string, nb::object> ExtractCompiledSqlAndParams(const nb::object &query, nb::object params);
379380
vector<unique_ptr<SQLStatement>> GetStatements(const nb::object &query);
380381

381382
static void DetectEnvironment();

‎src/pyconnection.cpp‎

Lines changed: 76 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -733,12 +733,75 @@ std::shared_ptr<DuckDBPyConnection> DuckDBPyConnection::ExecuteFromString(const
733733
return Execute(nb::str(query.c_str(), query.size()));
734734
}
735735

736+
pair<string, nb::object> DuckDBPyConnection::ExtractCompiledSqlAndParams(const nb::object &query, nb::object params) {
737+
nb::object compiled = query;
738+
if (!nb::hasattr(compiled, "sql") || !nb::hasattr(compiled, "params")) {
739+
if (!nb::hasattr(query, "compile")) {
740+
return {string(), params};
741+
}
742+
compiled = query.attr("compile")();
743+
}
744+
745+
if (!nb::hasattr(compiled, "sql") || !nb::hasattr(compiled, "params")) {
746+
return {string(), params};
747+
}
748+
749+
auto compiled_sql = nb::cast<std::string>(compiled.attr("sql"));
750+
auto compiled_params_obj = compiled.attr("params");
751+
if (!duckdb::PyUtil::IsDictLike(compiled_params_obj)) {
752+
throw InvalidInputException("Compiled SQL parameters must be a dictionary");
753+
}
754+
755+
auto compiled_params = nb::cast<nb::dict>(compiled_params_obj);
756+
if (compiled_params.size() == 0) {
757+
return {compiled_sql, params};
758+
}
759+
760+
if (params.is_none()) {
761+
return {compiled_sql, compiled_params};
762+
}
763+
764+
if (duckdb::PyUtil::IsDictLike(params)) {
765+
auto merged_params = nb::dict();
766+
for (auto item : compiled_params) { // nanobind dict iteration yields std::pair<handle,handle> by value
767+
merged_params[item.first] = item.second;
768+
}
769+
auto provided_params = nb::cast<nb::dict>(params);
770+
for (auto item : provided_params) { // nanobind dict iteration yields std::pair<handle,handle> by value
771+
if (merged_params.contains(item.first)) {
772+
throw nb::value_error("Cannot merge compiled SQL parameters with duplicate parameter names");
773+
}
774+
merged_params[item.first] = item.second;
775+
}
776+
return {compiled_sql, merged_params};
777+
}
778+
779+
if (duckdb::PyUtil::IsListLike(params)) {
780+
if (nb::len(params) == 0) {
781+
return {compiled_sql, compiled_params};
782+
}
783+
throw nb::value_error("Cannot merge compiled SQL named parameters with positional parameters");
784+
}
785+
786+
throw InvalidInputException("Prepared parameters can only be passed as a list or a dictionary");
787+
}
788+
736789
std::shared_ptr<DuckDBPyConnection> DuckDBPyConnection::Execute(const nb::object &query, nb::object params) {
737790
nb::gil_scoped_acquire gil;
738791
ConnectionLockGuard conn_lock(*this);
739792
con.SetResult(nullptr);
740793

741-
auto statements = GetStatements(query);
794+
auto normalized_query = ExtractCompiledSqlAndParams(query, params);
795+
auto &compiled_sql = normalized_query.first;
796+
auto &merged_params = normalized_query.second;
797+
vector<unique_ptr<SQLStatement>> statements;
798+
if (!compiled_sql.empty()) {
799+
statements = GetStatements(nb::str(compiled_sql.c_str(), compiled_sql.size()));
800+
params = merged_params;
801+
} else {
802+
statements = GetStatements(query);
803+
}
804+
742805
if (statements.empty()) {
743806
// TODO: should we throw?
744807
return nullptr;
@@ -1643,7 +1706,17 @@ std::unique_ptr<DuckDBPyRelation> DuckDBPyConnection::RunQuery(const nb::object
16431706
alias = "unnamed_relation_" + StringUtil::GenerateRandomName(16);
16441707
}
16451708

1646-
auto statements = GetStatements(query);
1709+
auto normalized_query = ExtractCompiledSqlAndParams(query, params);
1710+
auto &compiled_sql = normalized_query.first;
1711+
auto &merged_params = normalized_query.second;
1712+
vector<unique_ptr<SQLStatement>> statements;
1713+
if (!compiled_sql.empty()) {
1714+
statements = GetStatements(nb::str(compiled_sql.c_str(), compiled_sql.size()));
1715+
params = merged_params;
1716+
} else {
1717+
statements = GetStatements(query);
1718+
}
1719+
16471720
if (statements.empty()) {
16481721
// TODO: should we throw?
16491722
return nullptr;
@@ -1656,7 +1729,7 @@ std::unique_ptr<DuckDBPyRelation> DuckDBPyConnection::RunQuery(const nb::object
16561729

16571730
// Attempt to create a Relation for lazy execution if possible
16581731
shared_ptr<Relation> relation;
1659-
bool has_params = !nb::none().is(params) && nb::len(params) > 0;
1732+
bool has_params = !params.is_none() && nb::len(params) > 0;
16601733
if (!has_params) {
16611734
// No params (or empty params) — use lazy QueryRelation path
16621735
{

‎src/pyexpression/initialize.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,7 @@ void DuckDBPyExpression::Initialize(nb::module_ &m) {
338338
Print the stringified version of the expression.
339339
)";
340340
expression.def("show", &DuckDBPyExpression::Print, docs);
341+
expression.def("__duckdb_template__", &DuckDBPyExpression::ToString);
341342

342343
docs = R"(
343344
Set the order by modifier to ASCENDING.

‎src/pyrelation/initialize.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,7 @@ void DuckDBPyRelation::Initialize(nb::handle &m) {
366366
.def("show", &DuckDBPyRelation::Print, "Display a summary of the data", nb::kw_only(),
367367
nb::arg("max_width") = nb::none(), nb::arg("max_rows") = nb::none(), nb::arg("max_col_width") = nb::none(),
368368
nb::arg("null_value") = nb::none(), nb::arg("render_mode") = nb::none())
369+
.def("__duckdb_template__", &DuckDBPyRelation::ToSQL)
369370
.def("__str__", &DuckDBPyRelation::ToString)
370371
.def("__repr__", &DuckDBPyRelation::ToString);
371372

‎src/typing/pytype.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -356,6 +356,7 @@ void DuckDBPyType::Initialize(nb::handle &m) {
356356
auto type_module = nb::class_<DuckDBPyType>(m, "DuckDBPyType", nb::is_weak_referenceable());
357357

358358
type_module.def("__repr__", &DuckDBPyType::ToString, "Stringified representation of the type object");
359+
type_module.def("__duckdb_template__", &DuckDBPyType::ToString);
359360
type_module.def("__eq__", &DuckDBPyType::Equals, "Compare two types for equality", nb::arg("other"),
360361
nb::is_operator());
361362
type_module.def("__eq__", &DuckDBPyType::EqualsString, "Compare two types for equality", nb::arg("other"),

‎tests/fast/test_template_e2e.py‎

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,8 @@ def test_module_level_sql_apis_accept_sql_template() -> None:
4141

4242
def test_module_level_execute_accepts_sql_template() -> None:
4343
conn = duckdb.connect()
44-
query = template("SELECT ", "hello")
45-
assert duckdb.execute(query, connection=conn).fetchone() == ("hello",)
44+
query = template("SELECT ", 5)
45+
assert duckdb.execute(query, connection=conn).fetchone() == (5,)
4646

4747

4848
def test_connection_sql_accepts_alias_kwarg_with_template() -> None:
@@ -61,11 +61,19 @@ def test_connection_sql_template_can_merge_additional_params() -> None:
6161

6262
def test_connection_sql_template_param_name_conflict_with_additional_params_raises() -> None:
6363
conn = duckdb.connect()
64-
query = template("SELECT ", param(10, "num"), " + $num")
64+
query = template("SELECT ", param(10, "num", exact=True), " + $num")
6565
with pytest.raises((duckdb.InvalidInputException, ValueError)):
6666
conn.sql(query, params={"num": 5}).fetchall()
6767

6868

69+
def test_cant_merge_with_positional_params() -> None:
70+
conn = duckdb.connect()
71+
# It doesn't even have a name, but still should error
72+
query = template("SELECT ", 10, " + ?")
73+
with pytest.raises(ValueError, match="Cannot merge compiled SQL named parameters with positional parameters"):
74+
conn.sql(query, params=[5]).fetchall()
75+
76+
6977
def test_sql_apis_accept_compiled_sql() -> None:
7078
conn = duckdb.connect()
7179
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:
8391
assert conn.sql(query).fetchall() == [(0,), (2,), (4,)]
8492

8593

86-
def test_interpolated_strings_are_parameterized_by_default() -> None:
87-
conn = duckdb.connect()
88-
conn.execute("CREATE TABLE names(name VARCHAR)")
89-
conn.execute("INSERT INTO names VALUES ('alice'), ('bob')")
90-
91-
untrusted = "alice' OR 1=1 --"
92-
query = template("SELECT count(*) FROM names WHERE name = ", untrusted)
93-
assert conn.sql(query).fetchone() == (0,)
94-
95-
9694
def test_builtin_duckdbpytype_object_interpolates_in_template() -> None:
9795
conn = duckdb.connect()
9896
integer_type = duckdb.sqltype("INTEGER")

0 commit comments

Comments
 (0)