Skip to main content

cratestack_sqlx/query/write/
upsert.rs

1//! `INSERT … ON CONFLICT (<pk>) DO UPDATE …`, but with the
2//! create/update distinction *resolved* rather than merely guessed, so
3//! we can:
4//!
5//!   * pick the right policy slot (both must allow at call time)
6//!   * emit the correct ModelEventKind (Created vs Updated)
7//!   * capture an audit `before` snapshot only on the update branch
8//!
9//! A `SELECT … FOR UPDATE` probe inside the same transaction predicts
10//! the branch — binding when it finds a row (it holds the lock), a
11//! guess when it does not — so the statement itself gets the last word.
12//! See `upsert_resolve`'s module doc for how a mispredicted insert is
13//! recovered into a proper update (cratestack#745).
14//!
15//! The upsert is always transactional regardless of whether the model
16//! emits events or has `@@audit`. One extra round-trip for the
17//! SELECT, in exchange for clean event/audit semantics. Upsert is not
18//! a hot read path — callers who need raw insert/update throughput
19//! should use `.create()` / `.update()` directly.
20
21use cratestack_core::{CratestackContext, CratestackError};
22
23use crate::audit::{RunInTxOutcome, dispatch_audit_sink};
24use crate::{
25    ConflictTarget, ModelDescriptor, SqlxRuntime, UpsertModelInput, cratestack_error_from_sqlx,
26    sqlx,
27};
28
29use super::upsert_do_nothing::UpsertRecordDoNothing;
30use super::upsert_exec::run_upsert_in_tx;
31
32#[derive(Debug, Clone)]
33pub struct UpsertRecord<'a, M: 'static, PK: 'static, I> {
34    pub(crate) runtime: &'a SqlxRuntime,
35    pub(crate) descriptor: &'static ModelDescriptor<M, PK>,
36    pub(crate) input: I,
37    pub(crate) conflict_target: ConflictTarget,
38}
39
40impl<'a, M: 'static, PK: 'static, I> UpsertRecord<'a, M, PK, I>
41where
42    I: UpsertModelInput<M>,
43{
44    /// Choose the conflict target. Defaults to the model's primary
45    /// key; pass [`ConflictTarget::columns`] to upsert on a composite
46    /// unique key instead. The named columns must form a `UNIQUE`
47    /// constraint/index on the target table.
48    pub fn on_conflict(mut self, target: ConflictTarget) -> Self {
49        self.conflict_target = target;
50        self
51    }
52
53    /// Switch this call to `ON CONFLICT ... DO NOTHING` semantics
54    /// (cratestack#487), independent of `descriptor.upsert_update_columns`:
55    /// on conflict, leave the existing row completely untouched instead
56    /// of merging `upsert_update_columns` into it. This is the
57    /// idempotent-insert shape ledger-style writes need — e.g. a
58    /// cash-in claim that inserts a `PENDING` row and treats a
59    /// conflict as "already in flight" must never let a retry's blank
60    /// values overwrite an existing `COMPLETED` row's `transfer_ref`.
61    ///
62    /// Returns a distinct builder type rather than a flag on
63    /// `UpsertRecord` because the return shape genuinely changes: a
64    /// real `DO NOTHING` returns nothing on conflict, so the caller
65    /// needs `Inserted` vs `Existing` distinguishable in the type
66    /// ([`crate::UpsertOutcome`]) rather than collapsed into a plain
67    /// `M` the way `.run()` returns it today. Encoding that as a
68    /// separate type also means existing `.upsert(..).run(..)` callers
69    /// keep their `Result<M, CratestackError>` signature unchanged — this is
70    /// purely additive, not a behavior change for the DO UPDATE path.
71    pub fn do_nothing(self) -> UpsertRecordDoNothing<'a, M, PK, I> {
72        UpsertRecordDoNothing {
73            runtime: self.runtime,
74            descriptor: self.descriptor,
75            input: self.input,
76            conflict_target: self.conflict_target,
77        }
78    }
79
80    /// Render an approximate SQL preview. The actual upsert wraps a
81    /// `SELECT … FOR UPDATE` around the `INSERT … ON CONFLICT`, but
82    /// this preview returns only the conflict-bearing statement.
83    ///
84    /// Deliberately does NOT call [`ConflictTarget::validate`]
85    /// (cratestack#741 finding 3) — every other `preview_sql()` in
86    /// this codebase (`find_many`, `create`, `update`, `delete`, …)
87    /// returns a bare `String`, no `Result`, so it can never fail;
88    /// matching that established shape here means `cratestack-studio`'s
89    /// interactive SQL-preview tooling can keep calling it uniformly
90    /// across every builder without a special case for upsert. The
91    /// combination this would reject (a predicate on
92    /// [`ConflictTarget::PrimaryKey`]) is still caught before any SQL
93    /// *runs* — [`Self::run`]/[`Self::run_in_tx`] call `.validate()` via
94    /// `prepare_upsert_insert` first — so nothing unsafe executes;
95    /// only the preview string itself can show a `WHERE` clause paired
96    /// with a target `.validate()` would reject. A caller that wants to
97    /// know ahead of rendering can call `.validate()` on the same
98    /// `ConflictTarget` value passed to `.on_conflict(..)`.
99    pub fn preview_sql(&self) -> String {
100        let values = self.input.sql_values();
101        let placeholders = (1..=values.len())
102            .map(|index| format!("${index}"))
103            .collect::<Vec<_>>()
104            .join(", ");
105        let columns = values
106            .iter()
107            .map(|value| value.column)
108            .collect::<Vec<_>>()
109            .join(", ");
110        let update_assignments = self
111            .descriptor
112            .upsert_update_columns
113            .iter()
114            .map(|column| format!("{column} = EXCLUDED.{column}"))
115            .collect::<Vec<_>>()
116            .join(", ");
117        let version_bump = match self.descriptor.version_column {
118            Some(col) => format!(
119                ", {col} = {table}.{col} + 1",
120                table = self.descriptor.table_name,
121                col = col
122            ),
123            None => String::new(),
124        };
125        let conflict_tuple = match self.conflict_target.as_columns() {
126            None => self.descriptor.primary_key.to_owned(),
127            Some(cols) => cols.join(", "),
128        };
129        let conflict_predicate = match self.conflict_target.predicate() {
130            Some(predicate) => format!(" WHERE {predicate}"),
131            None => String::new(),
132        };
133
134        format!(
135            "INSERT INTO {table} ({columns}) VALUES ({placeholders}) \
136             ON CONFLICT ({conflict_tuple}){conflict_predicate} DO UPDATE SET {update_assignments}{version_bump} \
137             RETURNING {projection}",
138            table = self.descriptor.table_name,
139            projection = self.descriptor.select_projection(),
140        )
141    }
142
143    pub async fn run(self, ctx: &CratestackContext) -> Result<M, CratestackError>
144    where
145        for<'r> M: Send + Unpin + sqlx::FromRow<'r, sqlx::postgres::PgRow> + serde::Serialize,
146        PK: Send + sqlx::Type<sqlx::Postgres> + for<'q> sqlx::Encode<'q, sqlx::Postgres>,
147    {
148        let runtime = self.runtime;
149        let mut tx = runtime
150            .pool()
151            .begin()
152            .await
153            .map_err(cratestack_error_from_sqlx)?;
154        let (record, emits_event, audit_event) = run_upsert_in_tx(
155            &mut tx,
156            runtime,
157            self.descriptor,
158            self.input,
159            self.conflict_target,
160            ctx,
161        )
162        .await?;
163        tx.commit().await.map_err(cratestack_error_from_sqlx)?;
164        if emits_event {
165            let _ = runtime.drain_event_outbox().await;
166        }
167        if let Some(event) = &audit_event {
168            dispatch_audit_sink(runtime, std::slice::from_ref(event)).await;
169        }
170        Ok(record)
171    }
172
173    /// Like [`Self::run`] but participates in a caller-supplied
174    /// transaction. The conflict probe runs against `tx`, so the row
175    /// lock is held until the caller commits. Neither the event outbox
176    /// drain nor the `AuditSink` fan-out happens here — see
177    /// `create.rs`'s `run_in_tx` doc comment for the full contract and
178    /// how a caller opts into both after their own commit.
179    pub async fn run_in_tx<'tx>(
180        self,
181        tx: &mut sqlx::Transaction<'tx, sqlx::Postgres>,
182        ctx: &CratestackContext,
183    ) -> Result<RunInTxOutcome<M>, CratestackError>
184    where
185        for<'r> M: Send + Unpin + sqlx::FromRow<'r, sqlx::postgres::PgRow> + serde::Serialize,
186        PK: Send + sqlx::Type<sqlx::Postgres> + for<'q> sqlx::Encode<'q, sqlx::Postgres>,
187    {
188        let (record, _emits_event, audit_event) = run_upsert_in_tx(
189            tx,
190            self.runtime,
191            self.descriptor,
192            self.input,
193            self.conflict_target,
194            ctx,
195        )
196        .await?;
197        Ok(RunInTxOutcome::new(
198            record,
199            audit_event.into_iter().collect(),
200        ))
201    }
202}