cartel-pg 0.20.0

pub(in crate::api) mod stateful;

use crate::api::{sql, tx};
use crate::port::replies;
use crate::protocol::runtime::ready::Blocking;
use crate::{api, port, protocol, schema};
use api::task::{group, tx::Reserve};
use cartel_core::reply::recv::attach::Attach;
use o3::base::cell::region;
use port::requests;
use protocol::encode;

#[derive(Clone, Copy)]
pub(super) struct Begin {
    isolation: tx::Isolation,
    access: tx::AccessMode,
    deferral: tx::Deferral,
}

impl Begin {
    pub(super) const fn new(
        level: tx::Isolation,
        access: tx::AccessMode,
        deferral: tx::Deferral,
    ) -> Self {
        Self {
            isolation: level,
            access,
            deferral,
        }
    }
}

#[derive(Clone, Copy)]
enum Quoted {
    Savepoint,
    ReleaseSavepoint,
    RollbackToSavepoint,
    Listen,
    Unlisten,
}

#[derive(Clone, Copy)]
enum Sql<'a> {
    Raw(&'a str),
    Begin(Begin),
    StatementTimeout(sql::Timeout),
    Quoted(Quoted, &'a str),
}

impl Sql<'_> {
    fn write(self, out: &mut encode::Writer<'_>) -> Result<(), encode::Error> {
        match self {
            Self::Raw(sql) => {
                if let Some(position) = sql.bytes().position(|byte| byte == 0) {
                    return Err(encode::Error::StatementSqlNul { position });
                }
                out.extend(sql.as_bytes())?;
            }
            Self::Begin(begin) => {
                out.extend(b"BEGIN")?;
                match begin.isolation {
                    tx::Isolation::Default => {}
                    tx::Isolation::ReadCommitted => {
                        out.extend(b" ISOLATION LEVEL READ COMMITTED")?;
                    }
                    tx::Isolation::RepeatableRead => {
                        out.extend(b" ISOLATION LEVEL REPEATABLE READ")?;
                    }
                    tx::Isolation::Serializable => {
                        out.extend(b" ISOLATION LEVEL SERIALIZABLE")?;
                    }
                }
                match begin.access {
                    tx::AccessMode::Default => {}
                    tx::AccessMode::ReadWrite => out.extend(b" READ WRITE")?,
                    tx::AccessMode::ReadOnly => out.extend(b" READ ONLY")?,
                }
                match begin.deferral {
                    tx::Deferral::Default => {}
                    tx::Deferral::Deferred => out.extend(b" DEFERRABLE")?,
                }
            }
            Self::StatementTimeout(timeout) => {
                out.extend(b"SET LOCAL statement_timeout TO ")?;
                out.write_u32(timeout.milliseconds())?;
            }
            Self::Quoted(cmd, name) => {
                match cmd {
                    Quoted::Savepoint => out.extend(b"SAVEPOINT \"")?,
                    Quoted::ReleaseSavepoint => {
                        out.extend(b"RELEASE SAVEPOINT \"")?;
                    }
                    Quoted::RollbackToSavepoint => {
                        out.extend(b"ROLLBACK TO SAVEPOINT \"")?;
                    }
                    Quoted::Listen => out.extend(b"LISTEN \"")?,
                    Quoted::Unlisten => out.extend(b"UNLISTEN \"")?,
                }
                let mut parts = name.split('"');
                if let Some(first) = parts.next() {
                    out.extend(first.as_bytes())?;
                }
                for part in parts.by_ref() {
                    out.extend(b"\"\"")?;
                    out.extend(part.as_bytes())?;
                }
                out.push(b'"')?;
            }
        }
        Ok(())
    }
}

#[derive(Clone, Copy)]
pub(super) struct Request<'a> {
    sql: Sql<'a>,
    extra: Extra<'a>,
}

impl<'a> Request<'a> {
    pub(super) fn raw(sql: &'a str) -> Self {
        Self::raw_extra(sql, Extra::Plain)
    }

    pub(super) fn raw_extra(sql: &'a str, extra: Extra<'a>) -> Self {
        Self {
            sql: Sql::Raw(sql),
            extra,
        }
    }

    pub(super) fn begin(begin: Begin) -> Self {
        Self {
            sql: Sql::Begin(begin),
            extra: Extra::Plain,
        }
    }

    pub(super) fn timeout(timeout: sql::Timeout) -> Self {
        Self {
            sql: Sql::StatementTimeout(timeout),
            extra: Extra::Plain,
        }
    }

    pub(super) fn savepoint(name: &'a str) -> Self {
        Self::quoted(Quoted::Savepoint, name)
    }

    pub(super) fn release(name: &'a str) -> Self {
        Self::quoted(Quoted::ReleaseSavepoint, name)
    }

    pub(super) fn rollback(name: &'a str) -> Self {
        Self::quoted(Quoted::RollbackToSavepoint, name)
    }

    pub(super) fn listen(channel: &'a str) -> Self {
        Self::quoted(Quoted::Listen, channel)
    }

    pub(super) fn unlisten(channel: &'a str) -> Self {
        Self::quoted(Quoted::Unlisten, channel)
    }

    fn quoted(command: Quoted, name: &'a str) -> Self {
        Self {
            sql: Sql::Quoted(command, name),
            extra: Extra::Plain,
        }
    }

    fn write(self, out: &mut encode::Writer<'_>) -> Result<(), encode::Error> {
        out.parse_with("", |out| self.sql.write(out), &[])?;
        out.bind("", "", &[1], 0, &[1], |_| Ok(()))?;
        out.execute()?;
        match self.extra {
            Extra::Plain => out.sync()?,
            Extra::CopyIn { data } => {
                out.copy_data(data)?;
                out.copy_done()?;
                out.sync()?;
            }
            Extra::CopyInOpen => {}
        }
        Ok(())
    }
}

#[derive(Clone, Copy)]
pub(super) enum Extra<'a> {
    Plain,
    CopyIn { data: &'a [u8] },
    CopyInOpen,
}

/// One encoded request retaining its exact publication boundary and owner state.
pub(super) struct Pending<'d, S = ()> {
    pub(in crate::api) frame: requests::Complete<'d>,
    pub(in crate::api) boundary: requests::Boundary,
    pub(in crate::api) state: S,
}

impl<'d, S> Pending<'d, S> {
    pub(in crate::api) fn map_state<R>(self, map: impl FnOnce(S) -> R) -> Pending<'d, R> {
        Pending {
            frame: self.frame,
            boundary: self.boundary,
            state: map(self.state),
        }
    }
}

impl<'d> Pending<'d> {
    pub(in crate::api) fn attempt<'app, 'g, 'r, 's, 'b, 'c, B: ?Sized, I, const ID: u8>(
        self,
        client: api::Client<'app, 'g, 'r, 's, 'b, 'c, 'd, B, I, ID>,
        reply: &mut impl Attach<
            'app,
            'b,
            'c,
            'd,
            replies::Row<'d>,
            replies::ReplyFailure<'d>,
            replies::Tag<'g, 'r, 's, 'b, 'c, 'd>,
        >,
        region: &mut region::Token<region::Identity<'d>>,
    ) -> api::TxOutcome<'app, 'g, 'r, 's, 'b, 'c, 'd, B, I, ID>
    where
        I: schema::QuerySet + 'd,
        'd: 'app,
    {
        let req = self;
        if let Err(error) = api::check_open(client) {
            return api::TxOutcome::Failed(error);
        }
        let shared = &client.port.shared;
        let Some(hold) = Reserve::acquire(client, region) else {
            return match shared.transaction().blocking() {
                Blocking::Held(count) => {
                    api::TxOutcome::Failed(client.port.backpressure(count).into())
                }
                Blocking::Clear => api::TxOutcome::NoConn { request: req },
            };
        };
        let conn = hold.sink().connection();
        let api::Pending {
            frame,
            boundary,
            state: (),
        } = req;
        let ctx = requests::ReplyContext::new(conn, boundary, region);
        match client
            .port
            .request()
            .try_reply(frame, hold, reply, ctx, Reserve::commit)
        {
            Ok(lease) => api::TxOutcome::Enqueued { lease },
            Err((error, frame, hold)) => {
                let req = api::Pending {
                    frame,
                    boundary,
                    state: hold,
                };
                match api::Throttle::new(conn, error, req) {
                    Ok(throttle) => api::TxOutcome::Throttled {
                        throttle: throttle.release(),
                    },
                    Err(error) => api::TxOutcome::Failed(error.into()),
                }
            }
        }
    }

    pub(in crate::api) fn typed<Q: schema::TypedQuery, I: schema::QuerySet, const ID: u8>(
        port: &port::Port<'_, '_, '_, '_, '_, '_, 'd, I, ID>,
        params: Q::Params<'_>,
        boundary: api::Boundary,
    ) -> Result<Self, crate::Error> {
        if !schema::contains::<I, Q>() {
            return Err(crate::Error::QueryMetadata);
        }
        let frame = port.request().encode(|frame| {
            schema::QueryMeta::validate::<Q>()?.write(frame, params)?;
            match boundary {
                api::Boundary::Close => frame.sync(),
                api::Boundary::Open => Ok(()),
            }
        })?;
        let boundary = match boundary {
            api::Boundary::Close => requests::Boundary::Close,
            api::Boundary::Open => requests::Boundary::Open,
        };
        Ok(Self {
            frame,
            boundary,
            state: (),
        })
    }

    pub(in crate::api) fn group<'app, 'g, 'r, 's, 'b, 'c, Q, I, const ID: u8, const N: usize>(
        port: &'app port::Port<'app, 'g, 'r, 's, 'b, 'c, 'd, I, ID>,
        params: [Q::Params<'_>; N],
    ) -> Result<Pending<'d, group::Ready<'g, 'r, 's, 'b, 'c, 'd, N>>, crate::Error>
    where
        Q: schema::TypedQuery,
        I: schema::QuerySet,
        'd: 'app,
    {
        if !schema::contains::<I, Q>() {
            return Err(crate::Error::QueryMetadata);
        }
        let access = port
            .diagnostics
            .as_ref()
            .ok_or(crate::Error::GroupDisabled)?;
        let ready = group::Ready::new(access)?;
        let frame = port.request().encode(|frame| {
            let query = schema::QueryMeta::validate::<Q>()?;
            for params in params {
                query.write(frame, params)?;
            }
            frame.sync()
        })?;
        Ok(Pending {
            frame,
            boundary: requests::Boundary::Group,
            state: ready,
        })
    }

    pub(in crate::api) fn raw<I: schema::QuerySet, const ID: u8>(
        port: &port::Port<'_, '_, '_, '_, '_, '_, 'd, I, ID>,
        req: Request<'_>,
    ) -> Result<Self, crate::Error> {
        let boundary = match req.extra {
            Extra::Plain | Extra::CopyIn { .. } => requests::Boundary::Close,
            Extra::CopyInOpen => requests::Boundary::External,
        };
        let frame = port.request().encode(|frame| req.write(frame))?;
        Ok(Self {
            frame,
            boundary,
            state: (),
        })
    }

    pub(in crate::api) fn commit<I: schema::QuerySet, const ID: u8>(
        port: &port::Port<'_, '_, '_, '_, '_, '_, 'd, I, ID>,
    ) -> Result<Self, crate::Error> {
        let mut request = Self::raw(port, Request::raw("COMMIT"))?;
        request.boundary = requests::Boundary::Commit;
        Ok(request)
    }
}

impl<Q: schema::TypedQuery> schema::Admission<Q> {
    fn write(
        &self,
        out: &mut encode::Writer<'_>,
        params: Q::Params<'_>,
    ) -> Result<(), encode::Error> {
        use protocol::value::encode::BindWriter;
        out.bind(
            "",
            Q::STATEMENT_NAME,
            Q::PARAM_FORMAT_CODES,
            self.count(),
            Q::RESULT_FORMAT_CODES,
            |out| {
                let mut values = BindWriter::new(out, self);
                Q::encode(params, &mut values)?;
                values.finish()
            },
        )?;
        out.execute()
    }
}