#![allow(clippy::manual_async_fn)]
use crate::DjogiError;
use crate::context::DjogiContext;
use crate::expr::Expr;
use crate::expr::arithmetic::Numeric;
use crate::expr::node::ExprNode;
use crate::model::Model;
use crate::pg::accumulator::as_params;
use crate::pg::decode::FromPgRow;
use crate::query::field::FieldRef;
use crate::query::queryset::QuerySet;
use crate::query::sql::{build_insert_select, build_insert_select_returning};
use crate::query::terminal::auto_set_tenant;
use std::collections::HashSet;
use std::future::Future;
use std::marker::PhantomData;
#[must_use = "InsertSelectSource is lazy — drop one and the source projection is silently omitted"]
#[derive(Debug, Clone)]
pub struct InsertSelectSource<S: Model, V> {
pub(crate) node: ExprNode,
_phantom: PhantomData<fn() -> (S, V)>,
}
impl<S: Model, V> InsertSelectSource<S, V> {
#[must_use = "InsertSelectSource is lazy — drop one and the source projection is silently omitted"]
pub fn literal(v: V) -> Self
where
V: Into<Expr<V>>,
{
let expr: Expr<V> = v.into();
Self {
node: expr.node,
_phantom: PhantomData,
}
}
pub(crate) fn from_node(node: ExprNode) -> Self {
Self {
node,
_phantom: PhantomData,
}
}
}
impl<S: Model, V: Numeric> std::ops::Add for InsertSelectSource<S, V> {
type Output = InsertSelectSource<S, V>;
fn add(self, rhs: Self) -> Self::Output {
InsertSelectSource::from_node(ExprNode::Add(Box::new(self.node), Box::new(rhs.node)))
}
}
impl<S: Model, V: Numeric> std::ops::Sub for InsertSelectSource<S, V> {
type Output = InsertSelectSource<S, V>;
fn sub(self, rhs: Self) -> Self::Output {
InsertSelectSource::from_node(ExprNode::Sub(Box::new(self.node), Box::new(rhs.node)))
}
}
impl<S: Model, V: Numeric> std::ops::Mul for InsertSelectSource<S, V> {
type Output = InsertSelectSource<S, V>;
fn mul(self, rhs: Self) -> Self::Output {
InsertSelectSource::from_node(ExprNode::Mul(Box::new(self.node), Box::new(rhs.node)))
}
}
impl<S: Model, V: Numeric> std::ops::Div for InsertSelectSource<S, V> {
type Output = InsertSelectSource<S, V>;
fn div(self, rhs: Self) -> Self::Output {
InsertSelectSource::from_node(ExprNode::Div(Box::new(self.node), Box::new(rhs.node)))
}
}
impl<S: Model> std::ops::Add<InsertSelectSource<S, time::Duration>>
for InsertSelectSource<S, time::OffsetDateTime>
{
type Output = InsertSelectSource<S, time::OffsetDateTime>;
fn add(self, rhs: InsertSelectSource<S, time::Duration>) -> Self::Output {
InsertSelectSource::from_node(ExprNode::Add(Box::new(self.node), Box::new(rhs.node)))
}
}
impl<S: Model> std::ops::Sub<InsertSelectSource<S, time::Duration>>
for InsertSelectSource<S, time::OffsetDateTime>
{
type Output = InsertSelectSource<S, time::OffsetDateTime>;
fn sub(self, rhs: InsertSelectSource<S, time::Duration>) -> Self::Output {
InsertSelectSource::from_node(ExprNode::Sub(Box::new(self.node), Box::new(rhs.node)))
}
}
impl<S: Model, V> FieldRef<S, V> {
#[must_use = "InsertSelectSource is lazy — drop one and the source projection is silently omitted"]
pub fn as_insert_source(self) -> InsertSelectSource<S, V> {
InsertSelectSource::from_node(ExprNode::Field {
column: self.column(),
})
}
}
#[must_use = "column mappings are lazy — drop one and the INSERT silently omits the column"]
pub struct InsertSelectColumn<S: Model, T: Model> {
pub(crate) target_column: &'static str,
pub(crate) source: ExprNode,
_phantom: PhantomData<fn() -> (S, T)>,
}
impl<S: Model, T: Model> std::fmt::Debug for InsertSelectColumn<S, T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("InsertSelectColumn")
.field("source_table", &S::table_name())
.field("target_table", &T::table_name())
.field("target_column", &self.target_column)
.field("source", &self.source)
.finish()
}
}
impl<S: Model, T: Model> Clone for InsertSelectColumn<S, T> {
fn clone(&self) -> Self {
InsertSelectColumn {
target_column: self.target_column,
source: self.source.clone(),
_phantom: PhantomData,
}
}
}
impl<S: Model, T: Model> InsertSelectColumn<S, T> {
#[doc(hidden)]
pub fn target_column(&self) -> &'static str {
self.target_column
}
#[doc(hidden)]
pub(crate) fn source(&self) -> &ExprNode {
&self.source
}
}
impl<T: Model, V> FieldRef<T, V> {
#[must_use = "column mappings are lazy — drop one and the INSERT silently omits the column"]
pub fn copy_from<S: Model>(self, source: InsertSelectSource<S, V>) -> InsertSelectColumn<S, T> {
InsertSelectColumn {
target_column: self.column(),
source: source.node,
_phantom: PhantomData,
}
}
}
pub trait IntoInsertColumns<S: Model, T: Model> {
fn into_insert_columns(self) -> Vec<InsertSelectColumn<S, T>>;
}
impl<S: Model, T: Model> IntoInsertColumns<S, T> for InsertSelectColumn<S, T> {
fn into_insert_columns(self) -> Vec<InsertSelectColumn<S, T>> {
vec![self]
}
}
impl<S: Model, T: Model> IntoInsertColumns<S, T> for Vec<InsertSelectColumn<S, T>> {
fn into_insert_columns(self) -> Vec<InsertSelectColumn<S, T>> {
self
}
}
#[must_use = "InsertSelectStmt is inert — call .execute(ctx) to run the INSERT ... SELECT"]
pub struct InsertSelectStmt<S: Model, T: Model> {
pub(crate) source: QuerySet<S>,
pub(crate) columns: Vec<InsertSelectColumn<S, T>>,
pub(crate) _target: PhantomData<fn() -> T>,
}
impl<S: Model, T: Model> Clone for InsertSelectStmt<S, T> {
fn clone(&self) -> Self {
InsertSelectStmt {
source: self.source.clone(),
columns: self.columns.clone(),
_target: PhantomData,
}
}
}
impl<S: Model, T: Model> std::fmt::Debug for InsertSelectStmt<S, T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("InsertSelectStmt")
.field("source_table", &S::table_name())
.field("target_table", &T::table_name())
.field("source", &self.source)
.field("columns", &self.columns)
.finish()
}
}
impl<S: Model, T: Model> InsertSelectStmt<S, T> {
fn validate_execute(&self) -> Result<(), DjogiError> {
if self.columns.is_empty() {
return Err(DjogiError::Validation(format!(
"insert_into::<{}>: column mapping is empty; an INSERT...SELECT \
with no columns is invalid SQL. The closure passed to \
QuerySet::insert_into must return at least one column mapping \
via FieldRef::copy_from",
T::table_name(),
)));
}
let mut seen: HashSet<&'static str> = HashSet::with_capacity(self.columns.len());
for col in &self.columns {
if !seen.insert(col.target_column) {
return Err(DjogiError::Validation(format!(
"insert_into::<{}>: target column '{}' appears more than \
once in the column mapping; Postgres rejects duplicate \
columns in an INSERT column list (SQLSTATE 42701)",
T::table_name(),
col.target_column,
)));
}
}
if !self.source.prefetch_paths.is_empty() {
return Err(DjogiError::Validation(format!(
"insert_into::<{}>: source queryset has registered prefetch \
paths, which have no meaning for INSERT...SELECT (no rows \
are returned to the caller). Drop the .prefetch(...) calls \
before .insert_into(...)",
T::table_name(),
)));
}
if !self.source.select_related_paths.is_empty() {
return Err(DjogiError::Validation(format!(
"insert_into::<{}>: source queryset has registered \
select_related paths, which expand the SELECT list with \
aliased joined columns the INSERT...SELECT column-mapping \
closure cannot reference. Drop the .select_related(...) \
calls before .insert_into(...)",
T::table_name(),
)));
}
if self.source.cache_target.is_some() {
return Err(DjogiError::Validation(format!(
"insert_into::<{}>: source queryset is bound to a Punnu via \
.cache(...). INSERT...SELECT returns the affected row count, \
not rows, so the cache binding has nothing to insert. Drop \
the .cache(...) call before .insert_into(...)",
T::table_name(),
)));
}
if !matches!(self.source.lock, crate::query::lock::LockMode::None) {
return Err(DjogiError::Validation(format!(
"insert_into::<{}>: source queryset carries a row-level lock \
(FOR UPDATE / FOR SHARE / NOWAIT / SKIP LOCKED) which is \
not yet supported on INSERT...SELECT in djogi v0.1. Drop \
the .select_for_update() / .nowait() / .skip_locked() / \
.select_for_share() / .for_share_nowait() / \
.for_share_skip_locked() call before .insert_into(...); a \
follow-up issue can lift this restriction with an \
explicit opt-in",
T::table_name(),
)));
}
if !matches!(
self.source.distinct,
crate::query::queryset::DistinctMode::None,
) {
return Err(DjogiError::Validation(format!(
"insert_into::<{}>: source queryset carries .distinct() / \
.distinct_on(...) which is not yet supported on \
INSERT...SELECT in djogi v0.1. Drop the .distinct...() call \
before .insert_into(...); a follow-up issue can lift this \
restriction with an explicit opt-in",
T::table_name(),
)));
}
Ok(())
}
pub fn execute<'ctx>(
self,
ctx: &'ctx mut DjogiContext,
) -> impl Future<Output = Result<u64, DjogiError>> + Send + 'ctx
where
S: 'ctx,
T: 'ctx,
{
async move {
self.validate_execute()?;
if self.source.is_empty() {
return Ok(0);
}
auto_set_tenant::<T>(ctx).await?;
auto_set_tenant::<S>(ctx).await?;
let acc = build_insert_select::<S, T>(&self.source, &self.columns)
.map_err(DjogiError::from)?;
let (sql, binds) = acc.into_parts();
let params = as_params(&binds);
let rows_affected = ctx.execute(&sql, ¶ms).await?;
Ok(rows_affected)
}
}
pub fn execute_returning<'ctx>(
self,
ctx: &'ctx mut DjogiContext,
) -> impl Future<Output = Result<Vec<T>, DjogiError>> + Send + 'ctx
where
S: 'ctx,
T: 'ctx + FromPgRow,
{
async move {
self.validate_execute()?;
if self.source.is_empty() {
return Ok(Vec::new());
}
auto_set_tenant::<T>(ctx).await?;
auto_set_tenant::<S>(ctx).await?;
let acc = build_insert_select_returning::<S, T>(&self.source, &self.columns)
.map_err(DjogiError::from)?;
let (sql, binds) = acc.into_parts();
let params = as_params(&binds);
let rows = ctx.query_all(&sql, ¶ms).await?;
let mut results = Vec::with_capacity(rows.len());
for row in &rows {
results.push(T::from_pg_row(row)?);
}
Ok(results)
}
}
}
impl<S: Model> QuerySet<S> {
#[must_use = "InsertSelectStmt is inert — call .execute(ctx) to run the INSERT ... SELECT"]
pub fn insert_into<T, F, I>(self, f: F) -> InsertSelectStmt<S, T>
where
T: Model,
F: FnOnce(T::Fields, S::Fields) -> I,
I: IntoInsertColumns<S, T>,
{
let columns = f(T::Fields::default(), S::Fields::default()).into_insert_columns();
InsertSelectStmt {
source: self,
columns,
_target: PhantomData,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::descriptor::ModelDescriptor;
use crate::query::field::FieldRef;
struct Source;
impl crate::model::__sealed::Sealed for Source {}
#[allow(clippy::manual_async_fn)]
impl Model for Source {
type Pk = i64;
type Fields = ();
fn table_name() -> &'static str {
"sources"
}
fn pk_value(&self) -> &i64 {
unreachable!()
}
fn descriptor() -> &'static ModelDescriptor {
unreachable!()
}
fn get(
_ctx: &mut crate::context::DjogiContext,
_id: i64,
) -> impl std::future::Future<Output = Result<Self, crate::DjogiError>> + Send {
async { unreachable!() }
}
fn create(
_ctx: &mut crate::context::DjogiContext,
_v: Self,
) -> impl std::future::Future<Output = Result<Self, crate::DjogiError>> + Send {
async { unreachable!() }
}
fn save<'ctx>(
&'ctx mut self,
_ctx: &'ctx mut crate::context::DjogiContext,
) -> impl std::future::Future<Output = Result<(), crate::DjogiError>> + Send + 'ctx
{
async { unreachable!() }
}
fn delete(
self,
_ctx: &mut crate::context::DjogiContext,
) -> impl std::future::Future<Output = Result<(), crate::DjogiError>> + Send {
async { unreachable!() }
}
fn refresh_from_db<'ctx>(
&'ctx self,
_ctx: &'ctx mut crate::context::DjogiContext,
) -> impl std::future::Future<Output = Result<Self, crate::DjogiError>> + Send + 'ctx
{
async { unreachable!() }
}
}
struct Target;
impl crate::model::__sealed::Sealed for Target {}
#[allow(clippy::manual_async_fn)]
impl Model for Target {
type Pk = i64;
type Fields = ();
fn table_name() -> &'static str {
"targets"
}
fn pk_value(&self) -> &i64 {
unreachable!()
}
fn descriptor() -> &'static ModelDescriptor {
unreachable!()
}
fn get(
_ctx: &mut crate::context::DjogiContext,
_id: i64,
) -> impl std::future::Future<Output = Result<Self, crate::DjogiError>> + Send {
async { unreachable!() }
}
fn create(
_ctx: &mut crate::context::DjogiContext,
_v: Self,
) -> impl std::future::Future<Output = Result<Self, crate::DjogiError>> + Send {
async { unreachable!() }
}
fn save<'ctx>(
&'ctx mut self,
_ctx: &'ctx mut crate::context::DjogiContext,
) -> impl std::future::Future<Output = Result<(), crate::DjogiError>> + Send + 'ctx
{
async { unreachable!() }
}
fn delete(
self,
_ctx: &mut crate::context::DjogiContext,
) -> impl std::future::Future<Output = Result<(), crate::DjogiError>> + Send {
async { unreachable!() }
}
fn refresh_from_db<'ctx>(
&'ctx self,
_ctx: &'ctx mut crate::context::DjogiContext,
) -> impl std::future::Future<Output = Result<Self, crate::DjogiError>> + Send + 'ctx
{
async { unreachable!() }
}
}
#[test]
fn field_ref_copy_from_field_builds_column_mapping() {
let target_col: FieldRef<Target, i32> = FieldRef::new("view_count");
let source_col: FieldRef<Source, i32> = FieldRef::new("score");
let mapping: InsertSelectColumn<Source, Target> =
target_col.copy_from(source_col.as_insert_source());
assert_eq!(mapping.target_column(), "view_count");
assert!(matches!(
mapping.source(),
ExprNode::Field { column } if *column == "score"
));
}
#[test]
fn field_ref_copy_from_literal_builds_column_mapping() {
let target_col: FieldRef<Target, i32> = FieldRef::new("status_code");
let mapping: InsertSelectColumn<Source, Target> =
target_col.copy_from(InsertSelectSource::<Source, _>::literal(42i32));
assert_eq!(mapping.target_column(), "status_code");
assert!(matches!(mapping.source(), ExprNode::Literal(_)));
}
#[test]
fn into_insert_columns_single_wraps_in_vec() {
let target_col: FieldRef<Target, i32> = FieldRef::new("view_count");
let source_col: FieldRef<Source, i32> = FieldRef::new("score");
let mapping = target_col.copy_from(source_col.as_insert_source());
let v: Vec<InsertSelectColumn<Source, Target>> = mapping.into_insert_columns();
assert_eq!(v.len(), 1);
assert_eq!(v[0].target_column(), "view_count");
}
#[test]
fn into_insert_columns_vec_passes_through() {
let target_a: FieldRef<Target, i32> = FieldRef::new("view_count");
let target_b: FieldRef<Target, bool> = FieldRef::new("published");
let source_a: FieldRef<Source, i32> = FieldRef::new("score");
let source_b: FieldRef<Source, bool> = FieldRef::new("active");
let vs: Vec<InsertSelectColumn<Source, Target>> = vec![
target_a.copy_from(source_a.as_insert_source()),
target_b.copy_from(source_b.as_insert_source()),
];
let out = vs.into_insert_columns();
assert_eq!(out.len(), 2);
assert_eq!(out[0].target_column(), "view_count");
assert_eq!(out[1].target_column(), "published");
}
#[test]
fn insert_into_builds_stmt_with_columns() {
let qs: QuerySet<Source> = QuerySet::new();
let stmt = qs.insert_into::<Target, _, _>(|_t, _s| {
let target_col: FieldRef<Target, i32> = FieldRef::new("view_count");
let source_col: FieldRef<Source, i32> = FieldRef::new("score");
vec![target_col.copy_from(source_col.as_insert_source())]
});
assert_eq!(stmt.columns.len(), 1);
assert_eq!(stmt.columns[0].target_column(), "view_count");
}
#[test]
fn insert_select_stmt_clones_preserve_columns() {
let qs: QuerySet<Source> = QuerySet::new();
let stmt = qs.insert_into::<Target, _, _>(|_t, _s| {
let target_col: FieldRef<Target, i32> = FieldRef::new("view_count");
let source_col: FieldRef<Source, i32> = FieldRef::new("score");
vec![target_col.copy_from(source_col.as_insert_source())]
});
let cloned = stmt.clone();
assert_eq!(cloned.columns.len(), 1);
assert_eq!(cloned.columns[0].target_column(), "view_count");
}
#[test]
fn insert_select_stmt_single_mapping_via_into_insert_columns() {
let qs: QuerySet<Source> = QuerySet::new();
let stmt = qs.insert_into::<Target, _, _>(|_t, _s| {
let target_col: FieldRef<Target, i32> = FieldRef::new("view_count");
let source_col: FieldRef<Source, i32> = FieldRef::new("score");
target_col.copy_from(source_col.as_insert_source())
});
assert_eq!(stmt.columns.len(), 1);
assert_eq!(stmt.columns[0].target_column(), "view_count");
}
#[test]
fn insert_select_source_arithmetic_composes_same_source_tag() {
let source_col: FieldRef<Source, i32> = FieldRef::new("score");
let target_col: FieldRef<Target, i32> = FieldRef::new("view_count");
let composed: InsertSelectSource<Source, i32> =
source_col.as_insert_source() + InsertSelectSource::<Source, _>::literal(1i32);
let mapping: InsertSelectColumn<Source, Target> = target_col.copy_from(composed);
assert_eq!(mapping.target_column(), "view_count");
assert!(matches!(mapping.source(), ExprNode::Add(_, _)));
}
}