Skip to main content

cratestack_sqlx/query/write/
upsert.rs

1//! `INSERT … ON CONFLICT (<pk>) DO UPDATE …`, but with the
2//! create/update distinction made *before* the SQL runs (via a
3//! `SELECT … FOR UPDATE` probe inside the same transaction) so 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//! The upsert is always transactional regardless of whether the model
10//! emits events or has `@@audit`. One extra round-trip for the
11//! SELECT, in exchange for clean event/audit semantics. Upsert is not
12//! a hot read path — callers who need raw insert/update throughput
13//! should use `.create()` / `.update()` directly.
14
15use cratestack_core::{CoolContext, CoolError};
16
17use crate::{
18    ConflictTarget, ModelDescriptor, SqlxRuntime, UpsertModelInput, cool_error_from_sqlx, sqlx,
19};
20
21use super::upsert_exec::run_upsert_in_tx;
22
23#[derive(Debug, Clone)]
24pub struct UpsertRecord<'a, M: 'static, PK: 'static, I> {
25    pub(crate) runtime: &'a SqlxRuntime,
26    pub(crate) descriptor: &'static ModelDescriptor<M, PK>,
27    pub(crate) input: I,
28    pub(crate) conflict_target: ConflictTarget,
29}
30
31impl<'a, M: 'static, PK: 'static, I> UpsertRecord<'a, M, PK, I>
32where
33    I: UpsertModelInput<M>,
34{
35    /// Choose the conflict target. Defaults to the model's primary
36    /// key; pass [`ConflictTarget::Columns`] to upsert on a composite
37    /// unique key instead. The named columns must form a `UNIQUE`
38    /// constraint/index on the target table.
39    pub fn on_conflict(mut self, target: ConflictTarget) -> Self {
40        self.conflict_target = target;
41        self
42    }
43
44    /// Render an approximate SQL preview. The actual upsert wraps a
45    /// `SELECT … FOR UPDATE` around the `INSERT … ON CONFLICT`, but
46    /// this preview returns only the conflict-bearing statement.
47    pub fn preview_sql(&self) -> String {
48        let values = self.input.sql_values();
49        let placeholders = (1..=values.len())
50            .map(|index| format!("${index}"))
51            .collect::<Vec<_>>()
52            .join(", ");
53        let columns = values
54            .iter()
55            .map(|value| value.column)
56            .collect::<Vec<_>>()
57            .join(", ");
58        let update_assignments = self
59            .descriptor
60            .upsert_update_columns
61            .iter()
62            .map(|column| format!("{column} = EXCLUDED.{column}"))
63            .collect::<Vec<_>>()
64            .join(", ");
65        let version_bump = match self.descriptor.version_column {
66            Some(col) => format!(
67                ", {col} = {table}.{col} + 1",
68                table = self.descriptor.table_name,
69                col = col
70            ),
71            None => String::new(),
72        };
73        let conflict_tuple = match self.conflict_target {
74            ConflictTarget::PrimaryKey => self.descriptor.primary_key.to_owned(),
75            ConflictTarget::Columns(cols) => cols.join(", "),
76        };
77
78        format!(
79            "INSERT INTO {table} ({columns}) VALUES ({placeholders}) \
80             ON CONFLICT ({conflict_tuple}) DO UPDATE SET {update_assignments}{version_bump} \
81             RETURNING {projection}",
82            table = self.descriptor.table_name,
83            projection = self.descriptor.select_projection(),
84        )
85    }
86
87    pub async fn run(self, ctx: &CoolContext) -> Result<M, CoolError>
88    where
89        for<'r> M: Send + Unpin + sqlx::FromRow<'r, sqlx::postgres::PgRow> + serde::Serialize,
90        PK: Send + sqlx::Type<sqlx::Postgres> + for<'q> sqlx::Encode<'q, sqlx::Postgres>,
91    {
92        let runtime = self.runtime;
93        let mut tx = runtime.pool().begin().await.map_err(cool_error_from_sqlx)?;
94        let (record, emits_event) = run_upsert_in_tx(
95            &mut tx,
96            runtime,
97            self.descriptor,
98            self.input,
99            self.conflict_target,
100            ctx,
101        )
102        .await?;
103        tx.commit().await.map_err(cool_error_from_sqlx)?;
104        if emits_event {
105            let _ = runtime.drain_event_outbox().await;
106        }
107        Ok(record)
108    }
109
110    /// Like [`Self::run`] but participates in a caller-supplied
111    /// transaction. The conflict probe runs against `tx`, so the row
112    /// lock is held until the caller commits. The event outbox is not
113    /// drained here.
114    pub async fn run_in_tx<'tx>(
115        self,
116        tx: &mut sqlx::Transaction<'tx, sqlx::Postgres>,
117        ctx: &CoolContext,
118    ) -> Result<M, CoolError>
119    where
120        for<'r> M: Send + Unpin + sqlx::FromRow<'r, sqlx::postgres::PgRow> + serde::Serialize,
121        PK: Send + sqlx::Type<sqlx::Postgres> + for<'q> sqlx::Encode<'q, sqlx::Postgres>,
122    {
123        let (record, _) = run_upsert_in_tx(
124            tx,
125            self.runtime,
126            self.descriptor,
127            self.input,
128            self.conflict_target,
129            ctx,
130        )
131        .await?;
132        Ok(record)
133    }
134}