athena_rs 3.3.0

//! Shared query builder helpers for AthenaClient.
use crate::client::AthenaClient;
use crate::client::backend::{BackendError, QueryResult};
use crate::client::gateway_api::{
    GatewayDeleteRequest, GatewayFetchRequest, GatewayInsertRequest, GatewayRpcFilterOperator,
    GatewayRpcRequest, GatewayUpdateRequest,
};
use serde_json::Value;

/// ## `OrderDirection`
/// Direction used for ordering results.
///
/// # Arguments
///
/// * `Asc` - The ascending order.
/// * `Desc` - The descending order.
///
/// # Returns
///
/// A `OrderDirection` containing the order direction.
///
#[derive(Debug, Clone, Copy)]
pub enum OrderDirection {
    Asc,
    Desc,
}

/// ## `Condition`
/// Represents a column-based filter condition.
///
/// # Arguments
///
/// * `column` - The column.
/// * `operator` - The operator.
/// * `values` - The values.
///
/// # Returns
///
/// A `Condition` containing the condition.
///
#[derive(Debug, Clone)]
pub struct Condition {
    pub column: String,
    pub operator: ConditionOperator,
    pub values: Vec<Value>,
}

/// ## `ConditionOperator`
/// Supported comparison operators.
///
/// # Arguments
///
/// * `Eq` - The equal operator.
/// * `Neq` - The not equal operator.
/// * `Gt` - The greater than operator.
/// * `Lt` - The less than operator.
/// * `In` - The in operator.
///
/// # Returns
///
/// A `ConditionOperator` containing the condition operator.
///
#[derive(Debug, Clone, Copy)]
pub enum ConditionOperator {
    Eq,
    Neq,
    Gt,
    Lt,
    In,
}

impl Condition {
    /// ## `new`
    /// Create a new condition.
    ///
    /// # Arguments
    ///
    /// * `column` - The column.
    /// * `operator` - The operator.
    /// * `values` - The values.
    ///
    /// # Returns
    ///
    /// A `Condition` containing the condition.
    ///
    pub fn new(column: impl Into<String>, operator: ConditionOperator, values: Vec<Value>) -> Self {
        Self {
            column: column.into(),
            operator,
            values,
        }
    }
}

/// Select query builder.
///
/// # Arguments
///
/// * `client` - The client.
/// * `table` - The table.
/// * `columns` - The columns.
/// * `raw_select` - The raw select.
/// * `conditions` - The conditions.
/// * `order_by` - The order by.
/// * `limit` - The limit.
/// * `offset` - The offset.
///
/// # Returns
///
/// A `SelectBuilder` containing the select builder.
///
pub struct SelectBuilder<'a> {
    client: &'a AthenaClient,
    table: String,
    columns: Vec<String>,
    raw_select: Option<String>,
    conditions: Vec<Condition>,
    order_by: Vec<(String, OrderDirection)>,
    limit: Option<usize>,
    offset: Option<usize>,
}

impl<'a> SelectBuilder<'a> {
    /// ## `new`
    /// Create a new select builder.
    ///
    /// # Arguments
    ///
    /// * `client` - The client.
    /// * `table` - The table.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub(crate) fn new(client: &'a AthenaClient, table: impl Into<String>) -> Self {
        Self {
            client,
            table: table.into(),
            columns: Vec::new(),
            raw_select: None,
            conditions: Vec::new(),
            order_by: Vec::new(),
            limit: None,
            offset: None,
        }
    }

    /// ## `columns`
    /// Set the columns.
    ///
    /// # Arguments
    ///
    /// * `columns` - The columns.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn columns(mut self, columns: impl IntoIterator<Item = impl Into<String>>) -> Self {
        self.columns = columns.into_iter().map(|c| c.into()).collect();
        self
    }

    /// ## `raw_select`
    /// Provide a Supabase/PostgREST-style `select` string with nested relations,
    /// e.g. `id,name,instruments(id,name)`.
    ///
    /// # Arguments
    ///
    /// * `select` - The raw select.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn raw_select(mut self, select: impl Into<String>) -> Self {
        self.raw_select = Some(select.into());
        self
    }

    /// ## `add_condition`
    /// Add a condition.
    ///
    /// # Arguments
    ///
    /// * `column` - The column.
    /// * `operator` - The operator.
    /// * `values` - The values.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    fn add_condition(
        mut self,
        column: &str,
        operator: ConditionOperator,
        values: Vec<Value>,
    ) -> Self {
        self.conditions
            .push(Condition::new(column, operator, values));
        self
    }

    /// Add a pre-built condition.
    pub fn where_condition(mut self, condition: Condition) -> Self {
        self.conditions.push(condition);
        self
    }

    /// Add multiple pre-built conditions.
    pub fn where_conditions(mut self, conditions: impl IntoIterator<Item = Condition>) -> Self {
        self.conditions.extend(conditions);
        self
    }

    /// ## `where_eq`
    /// Add an equality condition.
    ///
    /// # Arguments
    ///
    /// * `column` - The column.
    /// * `value` - The value.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn where_eq(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Eq, vec![value.into()])
    }

    /// ## `where_neq`
    /// Add a not equal condition.
    ///
    /// # Arguments
    ///
    /// * `column` - The column.
    /// * `value` - The value.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn where_neq(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Neq, vec![value.into()])
    }

    /// ## `where_gt`
    /// Add a greater than condition.
    ///
    /// # Arguments
    ///
    /// * `column` - The column.
    /// * `value` - The value.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn where_gt(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Gt, vec![value.into()])
    }

    /// ## `where_lt`
    /// Add a less than condition.
    ///
    /// # Arguments
    ///
    /// * `column` - The column.
    /// * `value` - The value.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn where_lt(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Lt, vec![value.into()])
    }

    /// ## `where_in`
    /// Add an in condition.
    ///
    /// # Arguments
    ///
    /// * `column` - The column.
    /// * `values` - The values.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn where_in(self, column: &str, values: Vec<Value>) -> Self {
        self.add_condition(column, ConditionOperator::In, values)
    }

    /// ## `order_by`
    /// Add an order by condition.
    ///
    /// # Arguments
    ///
    /// * `column` - The column.
    /// * `direction` - The direction.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn order_by(mut self, column: &str, direction: OrderDirection) -> Self {
        self.order_by.push((column.to_string(), direction));
        self
    }

    /// ## `limit`
    /// Set the limit.
    ///
    /// # Arguments
    ///
    /// * `value` - The limit.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn limit(mut self, value: usize) -> Self {
        self.limit = Some(value);
        self
    }

    /// ## `offset`
    /// Set the offset.
    ///
    /// # Arguments
    ///
    /// * `value` - The offset.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the select builder.
    ///
    pub fn offset(mut self, value: usize) -> Self {
        self.offset = Some(value);
        self
    }

    /// ## `execute`
    /// Execute the select query.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    pub async fn execute(self) -> Result<QueryResult, BackendError> {
        self.client.execute_select(self).await
    }

    /// ## `into_request`
    /// Convert the select builder into a typed gateway fetch request.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `GatewayFetchRequest` containing the select query.
    ///
    pub(crate) fn into_request(self) -> GatewayFetchRequest {
        GatewayFetchRequest {
            table: self.table,
            columns: self.columns,
            raw_select: self.raw_select,
            conditions: self.conditions,
            order_by: self.order_by,
            limit: self.limit,
            offset: self.offset,
        }
    }
}

/// ## `InsertBuilder`
/// Insert query builder.
///
/// # Arguments
///
/// * `client` - The client.
/// * `table` - The table.
/// * `payload` - The payload.
///
/// # Returns
///
/// A `InsertBuilder` containing the insert builder.
///
pub struct InsertBuilder<'a> {
    client: &'a AthenaClient,
    table: String,
    payload: Value,
}

impl<'a> InsertBuilder<'a> {
    /// ## `new`
    /// Create a new insert builder.
    ///
    /// # Arguments
    ///
    /// * `client` - The client.
    /// * `table` - The table.
    ///
    /// # Returns
    ///
    /// A `InsertBuilder` containing the insert builder.
    ///
    pub(crate) fn new(client: &'a AthenaClient, table: impl Into<String>) -> Self {
        Self {
            client,
            table: table.into(),
            payload: Value::Null,
        }
    }

    /// ## `payload`
    /// Set the payload.
    ///
    /// # Arguments
    ///
    /// * `payload` - The payload.
    ///
    /// # Returns
    ///
    /// A `InsertBuilder` containing the insert builder.
    ///
    pub fn payload(mut self, payload: Value) -> Self {
        self.payload = payload;
        self
    }

    /// ## `execute`
    /// Execute the insert query.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    pub async fn execute(self) -> Result<QueryResult, BackendError> {
        self.client.execute_insert(self).await
    }

    /// ## `into_request`
    /// Convert the insert builder into a typed gateway insert request.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `GatewayInsertRequest` containing the insert query.
    ///
    pub(crate) fn into_request(self) -> GatewayInsertRequest {
        GatewayInsertRequest::new(self.table, self.payload)
    }
}

/// ## `UpdateBuilder`
/// Update query builder.
///
/// # Arguments
///
/// * `client` - The client.
/// * `table` - The table.
/// * `row_id` - The row id.
/// * `payload` - The payload.
///
/// # Returns
///
/// A `UpdateBuilder` containing the update builder.
///
pub struct UpdateBuilder<'a> {
    client: &'a AthenaClient,
    table: String,
    row_id: Option<String>,
    conditions: Vec<Condition>,
    payload: Value,
    allow_unfiltered: bool,
}

impl<'a> UpdateBuilder<'a> {
    /// ## `new`
    /// Create a new update builder.
    ///
    /// # Arguments
    ///
    /// * `client` - The client.
    /// * `table` - The table.
    /// * `row_id` - The row id.
    ///
    /// # Returns
    ///
    /// A `UpdateBuilder` containing the update builder.
    ///
    pub(crate) fn new(
        client: &'a AthenaClient,
        table: impl Into<String>,
        row_id: Option<String>,
    ) -> Self {
        Self {
            client,
            table: table.into(),
            row_id,
            conditions: Vec::new(),
            payload: Value::Null,
            allow_unfiltered: false,
        }
    }

    /// ## `row_id`
    /// Set the row id filter (`id = <row_id>`).
    pub fn row_id(mut self, row_id: impl Into<String>) -> Self {
        self.row_id = Some(row_id.into());
        self
    }

    fn add_condition(
        mut self,
        column: &str,
        operator: ConditionOperator,
        values: Vec<Value>,
    ) -> Self {
        self.conditions
            .push(Condition::new(column, operator, values));
        self
    }

    /// Add a pre-built condition.
    pub fn where_condition(mut self, condition: Condition) -> Self {
        self.conditions.push(condition);
        self
    }

    /// Add multiple pre-built conditions.
    pub fn where_conditions(mut self, conditions: impl IntoIterator<Item = Condition>) -> Self {
        self.conditions.extend(conditions);
        self
    }

    /// ## `where_eq`
    /// Add an equality condition.
    pub fn where_eq(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Eq, vec![value.into()])
    }

    /// ## `where_neq`
    /// Add a not equal condition.
    pub fn where_neq(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Neq, vec![value.into()])
    }

    /// ## `where_gt`
    /// Add a greater than condition.
    pub fn where_gt(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Gt, vec![value.into()])
    }

    /// ## `where_lt`
    /// Add a less than condition.
    pub fn where_lt(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Lt, vec![value.into()])
    }

    /// ## `where_in`
    /// Add an IN condition.
    pub fn where_in(self, column: &str, values: Vec<Value>) -> Self {
        self.add_condition(column, ConditionOperator::In, values)
    }

    /// ## `payload`
    /// Set the payload.
    ///
    /// # Arguments
    ///
    /// * `payload` - The payload.
    ///
    /// # Returns
    ///
    /// A `UpdateBuilder` containing the update builder.
    ///
    pub fn payload(mut self, payload: Value) -> Self {
        self.payload = payload;
        self
    }

    /// ## `unsafe_unfiltered`
    /// Explicitly opt-in to executing an update without any `row_id`/`where_*` filters.
    ///
    /// This should only be used for intentional full-table updates.
    pub fn unsafe_unfiltered(mut self) -> Self {
        self.allow_unfiltered = true;
        self
    }

    /// ## `execute`
    /// Execute the update query.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    pub async fn execute(self) -> Result<QueryResult, BackendError> {
        self.client.execute_update(self).await
    }

    /// ## `into_request`
    /// Convert the update builder into a typed gateway update request.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `GatewayUpdateRequest` containing the update query.
    ///
    pub(crate) fn into_request(self) -> GatewayUpdateRequest {
        let mut request = GatewayUpdateRequest::new(self.table, self.payload);
        if let Some(row_id) = self.row_id {
            request = request.row_id(row_id);
        }
        request = request.where_conditions(self.conditions);
        if self.allow_unfiltered {
            request = request.unsafe_unfiltered();
        }
        request
    }
}

/// Delete query builder.
///
/// # Arguments
///
/// * `client` - The client.
/// * `table` - The table.
/// * `row_id` - The row id.
///
/// # Returns
///
/// A `DeleteBuilder` containing the delete builder.
///
pub struct DeleteBuilder<'a> {
    client: &'a AthenaClient,
    table: String,
    row_id: Option<String>,
    conditions: Vec<Condition>,
    allow_unfiltered: bool,
}

impl<'a> DeleteBuilder<'a> {
    /// ## `new`
    /// Create a new delete builder.
    ///
    /// # Arguments
    ///
    /// * `client` - The client.
    /// * `table` - The table.
    /// * `row_id` - The row id.
    ///
    /// # Returns
    ///
    /// A `DeleteBuilder` containing the delete builder.
    ///
    pub(crate) fn new(
        client: &'a AthenaClient,
        table: impl Into<String>,
        row_id: Option<String>,
    ) -> Self {
        Self {
            client,
            table: table.into(),
            row_id,
            conditions: Vec::new(),
            allow_unfiltered: false,
        }
    }

    /// ## `row_id`
    /// Set the row id filter (`id = <row_id>`).
    pub fn row_id(mut self, row_id: impl Into<String>) -> Self {
        self.row_id = Some(row_id.into());
        self
    }

    fn add_condition(
        mut self,
        column: &str,
        operator: ConditionOperator,
        values: Vec<Value>,
    ) -> Self {
        self.conditions
            .push(Condition::new(column, operator, values));
        self
    }

    /// Add a pre-built condition.
    pub fn where_condition(mut self, condition: Condition) -> Self {
        self.conditions.push(condition);
        self
    }

    /// Add multiple pre-built conditions.
    pub fn where_conditions(mut self, conditions: impl IntoIterator<Item = Condition>) -> Self {
        self.conditions.extend(conditions);
        self
    }

    /// ## `where_eq`
    /// Add an equality condition.
    pub fn where_eq(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Eq, vec![value.into()])
    }

    /// ## `where_neq`
    /// Add a not equal condition.
    pub fn where_neq(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Neq, vec![value.into()])
    }

    /// ## `where_gt`
    /// Add a greater than condition.
    pub fn where_gt(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Gt, vec![value.into()])
    }

    /// ## `where_lt`
    /// Add a less than condition.
    pub fn where_lt(self, column: &str, value: impl Into<Value>) -> Self {
        self.add_condition(column, ConditionOperator::Lt, vec![value.into()])
    }

    /// ## `where_in`
    /// Add an IN condition.
    pub fn where_in(self, column: &str, values: Vec<Value>) -> Self {
        self.add_condition(column, ConditionOperator::In, values)
    }

    /// ## `unsafe_unfiltered`
    /// Explicitly opt-in to executing a delete without any `row_id`/`where_*` filters.
    ///
    /// This should only be used for intentional full-table deletes.
    pub fn unsafe_unfiltered(mut self) -> Self {
        self.allow_unfiltered = true;
        self
    }

    /// ## `execute`
    /// Execute the delete query.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    pub async fn execute(self) -> Result<QueryResult, BackendError> {
        self.client.execute_delete(self).await
    }

    /// ## `into_request`
    /// Convert the delete builder into a typed gateway delete request.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `GatewayDeleteRequest` containing the delete query.
    ///
    pub(crate) fn into_request(self) -> GatewayDeleteRequest {
        GatewayDeleteRequest {
            table: self.table,
            row_id: self.row_id,
            conditions: self.conditions,
            allow_unfiltered: self.allow_unfiltered,
        }
    }
}

/// RPC query builder.
pub struct RpcBuilder<'a> {
    client: &'a AthenaClient,
    request: GatewayRpcRequest,
}

impl<'a> RpcBuilder<'a> {
    pub(crate) fn new(client: &'a AthenaClient, function: impl Into<String>, args: Value) -> Self {
        Self {
            client,
            request: GatewayRpcRequest::new(function, args),
        }
    }

    pub fn schema(mut self, schema: impl Into<String>) -> Self {
        self.request = self.request.schema(schema);
        self
    }

    pub fn select(mut self, select: impl Into<String>) -> Self {
        self.request = self.request.select(select);
        self
    }

    fn add_filter(
        mut self,
        column: impl Into<String>,
        operator: GatewayRpcFilterOperator,
        value: Value,
    ) -> Self {
        self.request = self.request.filter(column, operator, value);
        self
    }

    pub fn eq(self, column: impl Into<String>, value: impl Into<Value>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::Eq, value.into())
    }

    pub fn neq(self, column: impl Into<String>, value: impl Into<Value>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::Neq, value.into())
    }

    pub fn gt(self, column: impl Into<String>, value: impl Into<Value>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::Gt, value.into())
    }

    pub fn gte(self, column: impl Into<String>, value: impl Into<Value>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::Gte, value.into())
    }

    pub fn lt(self, column: impl Into<String>, value: impl Into<Value>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::Lt, value.into())
    }

    pub fn lte(self, column: impl Into<String>, value: impl Into<Value>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::Lte, value.into())
    }

    pub fn in_(self, column: impl Into<String>, values: Vec<Value>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::In, Value::Array(values))
    }

    pub fn like(self, column: impl Into<String>, value: impl Into<Value>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::Like, value.into())
    }

    pub fn ilike(self, column: impl Into<String>, value: impl Into<Value>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::ILike, value.into())
    }

    pub fn is_null(self, column: impl Into<String>) -> Self {
        self.add_filter(column, GatewayRpcFilterOperator::Is, Value::Null)
    }

    pub fn order(mut self, column: impl Into<String>, direction: OrderDirection) -> Self {
        self.request = self.request.order_by(column, direction);
        self
    }

    pub fn limit(mut self, value: usize) -> Self {
        self.request = self.request.limit(value);
        self
    }

    pub fn offset(mut self, value: usize) -> Self {
        self.request = self.request.offset(value);
        self
    }

    pub fn count_exact(mut self) -> Self {
        self.request = self.request.count_exact();
        self
    }

    pub async fn execute(self) -> Result<QueryResult, BackendError> {
        self.client.execute_rpc(self).await
    }

    pub(crate) fn into_request(self) -> GatewayRpcRequest {
        self.request
    }
}