1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
use std::sync::Arc;
use async_trait::async_trait;
use dataflow_rs::engine::task_context::TaskContext;
use super::connector_handler::{ConnectorHandler, Produced};
use super::connector_helpers::{
ConnectorCall, QueryBudget, QueryFailure, acquire_conn, encode_failure, reject_mongo_connector,
require_op_allowed, to_connect_error,
};
use super::db_read::DbRead;
use super::schema::{FieldKind, FieldSchema};
use super::templated_input::TemplatedInput;
use crate::connector::ConnectorRegistry;
use crate::connector::pool_cache::SqlPoolCache;
use crate::engine::HandlerError;
/// Executes SQL write queries (INSERT, UPDATE, DELETE) against external databases
/// configured via connectors.
pub struct DbWriteHandler {
pub pool_cache: Arc<SqlPoolCache>,
pub registry: Arc<ConnectorRegistry>,
}
#[async_trait]
impl ConnectorHandler for DbWriteHandler {
const NAME: &'static str = "db_write";
type Kind = crate::connector::kind::Db;
type Input = TemplatedInput;
/// The same shape `db_read` parses — a literal statement and message-derived
/// binds — because the two differ in what the database does with it, not in
/// what the task says.
type Parsed = DbRead;
fn registry(&self) -> &Arc<ConnectorRegistry> {
&self.registry
}
fn parse(
&self,
call: &ConnectorCall<'_>,
input: &TemplatedInput,
ctx: &TaskContext<'_>,
) -> Result<Self::Parsed, HandlerError> {
DbRead::parse_statement(call, input, ctx)
}
fn gate(
_parsed: &Self::Parsed,
conn: &crate::connector::DbConnectorConfig,
connector: &str,
) -> Result<(), HandlerError> {
// Raw SQL cannot be classified per-op; it has its own gate.
require_op_allowed(&conn.operations, "raw_write", connector)?;
Ok(reject_mongo_connector(
<Self as ConnectorHandler>::NAME,
connector,
conn,
)?)
}
async fn run(
&self,
write: Self::Parsed,
db_config: &crate::connector::DbConnectorConfig,
call: &ConnectorCall<'_>,
_input: &TemplatedInput,
_ctx: &mut TaskContext<'_>,
) -> Result<Produced, HandlerError> {
let pool = self
.pool_cache
.get_pool(call.connector, db_config)
.await
.map_err(to_connect_error)?;
let query = write.query();
let params = write.params();
// One budget over both legs, so naming the connection does not turn
// the connector's `query_timeout_ms` into `connect_timeout_ms` plus it.
let budget = QueryBudget::start(db_config.query_timeout_ms);
let scalars: Vec<crate::connector::sql_encode::Scalar> =
params.iter().map(Into::into).collect();
let (rows_affected, last_insert_id) = crate::connector::pool_cache::dispatch_sql_pool!(
&pool, p, _decode, bind, typed_args, write_result => {
// Named for the same reason as `db_read`'s: a prepared
// statement's parameter types are cached per connection, so
// the prepare and the execute have to share one.
let mut conn = acquire_conn(&budget, call.name, p).await?;
let bound = budget
.run(call.name, async {
typed_args(&mut conn, query, Some(&scalars))
.await
.map_err(|e| QueryFailure::Classified(encode_failure("db_write", e)))
})
.await?;
// `AssertSqlSafe` states what this handler is: the
// raw-SQL escape hatch, whose statement is authored in the
// workflow rather than assembled here. sqlx 0.9 asks the
// caller to own that, and the answer is the same as it was
// before it asked — the text comes from the definition, the
// author-supplied *values* travel as bind parameters beside
// it, and a definition reaches the runtime only through the
// admin API. Authors who want values checked use the
// portable dialect (`data_query`/`data_write`) instead.
let sqlx_query = match bound {
crate::connector::sql_encode::Bound::Typed(args) => {
sqlx::query_with(sqlx::AssertSqlSafe(query), args)
}
crate::connector::sql_encode::Bound::Fallback { cache } => {
bind(sqlx::query(sqlx::AssertSqlSafe(query)), params).persistent(cache)
}
};
let result = budget
.run(call.name, sqlx_query.execute(&mut *conn))
.await?;
write_result(&result)
}
);
let mut out = serde_json::json!({ "rows_affected": rows_affected });
// The generated key, on the drivers that report one — the same answer
// `data_write` gives, so reaching for the escape hatch does not cost
// the most common follow-up question after an insert.
//
// Only for an insert, and that is not pedantry: SQLite's
// `last_insert_rowid` is a property of the *connection*, so after an
// UPDATE it reports whatever an earlier insert on that pooled
// connection left behind. Raw SQL cannot be classified per-op — which
// is why `db_write` has its own `raw_write` gate — but the leading
// keyword is enough to know whether an id was generated at all.
if let Some(id) = last_insert_id
&& matches!(
super::db_read::leading_keyword(query).as_deref(),
Some("INSERT" | "REPLACE")
)
{
out["last_insert_id"] = serde_json::json!(id);
}
Ok(out.into())
}
}
// -- Input schema (F53) --
//
// The table describing this handler's `function.input` lives next to the
// handler it describes. It used to sit in `schema.rs` with the other nine,
// which is how every schema/handler divergence in the 1.0 audit happened:
// a field was added, renamed or made conditional here and the table saying
// so was in a different file.
pub(super) const DB_WRITE_FIELDS: &[FieldSchema] = &[
FieldSchema {
name: "connector",
description: "Name of the SQL connector to execute against.",
kind: FieldKind::String,
required: true,
..FieldSchema::DEFAULT
},
FieldSchema {
name: "query",
description: "INSERT/UPDATE/DELETE statement. Bind placeholders are the backend's own spelling: ? for SQLite and MySQL, $1, $2, ... for PostgreSQL.",
kind: FieldKind::String,
required: true,
..FieldSchema::DEFAULT
},
FieldSchema {
name: "params",
description: "Array of values to bind to query placeholders, in order. Accepts {\"var\": \"path\"} to read the value from the message.",
kind: FieldKind::Array,
resolvable: true,
..FieldSchema::DEFAULT
},
FieldSchema {
name: "output",
description: "Dotted path where the write result is written: {\"rows_affected\": n}, plus \"last_insert_id\" for an INSERT on MySQL or SQLite (PostgreSQL reports generated columns through RETURNING instead).",
kind: FieldKind::String,
template_at: &[""],
..FieldSchema::DEFAULT
},
];