picodata-driver 3.0.0

use crate::metadata::DkMeta;

/// OID constants for supported PG types.
mod oid {
    pub const BOOL: u32 = 16;
    pub const INT2: u32 = 21;
    pub const INT4: u32 = 23;
    pub const INT8: u32 = 20;
    pub const FLOAT4: u32 = 700;
    pub const FLOAT8: u32 = 701;
    pub const TEXT: u32 = 25;
    pub const VARCHAR: u32 = 1043;
    pub const TIMESTAMP: u32 = 1114;
    pub const TIMESTAMPTZ: u32 = 1184;
    pub const UUID: u32 = 2950;
}

/// Encodes a single parameter as msgpack bytes. Returns `false` if the OID is unsupported.
fn encode_param_msgpack(
    param: &(dyn pg::types::ToSql + Sync),
    type_oid: u32,
    out: &mut Vec<u8>,
) -> bool {
    let Some(pg_type) = pg_type_from_oid(type_oid) else {
        log::warn!("encode_param_msgpack: unsupported OID {type_oid}");
        return false;
    };
    let mut pg_buf = bytes::BytesMut::with_capacity(32);
    let Ok(is_null) = param.to_sql_checked(&pg_type, &mut pg_buf) else {
        return false;
    };
    if matches!(is_null, pg::types::IsNull::Yes) {
        rmp::encode::write_nil(out).expect("write to Vec never fails");
        return true;
    }

    match type_oid {
        oid::BOOL => {
            let v = pg_buf.first().is_some_and(|&b| b != 0);
            rmp::encode::write_bool(out, v).expect("write to Vec never fails");
        }
        oid::INT2 => {
            let v = i16::from_be_bytes([pg_buf[0], pg_buf[1]]);
            rmp::encode::write_sint(out, i64::from(v)).expect("write to Vec never fails");
        }
        oid::INT4 => {
            let v = i32::from_be_bytes([pg_buf[0], pg_buf[1], pg_buf[2], pg_buf[3]]);
            rmp::encode::write_sint(out, i64::from(v)).expect("write to Vec never fails");
        }
        oid::INT8 => {
            let v = i64::from_be_bytes(pg_buf[..8].try_into().expect("pg INT8 is always 8 bytes"));
            rmp::encode::write_sint(out, v).expect("write to Vec never fails");
        }
        oid::FLOAT4 => {
            let v =
                f32::from_be_bytes(pg_buf[..4].try_into().expect("pg FLOAT4 is always 4 bytes"));
            rmp::encode::write_f32(out, v).expect("write to Vec never fails");
        }
        oid::FLOAT8 => {
            let v =
                f64::from_be_bytes(pg_buf[..8].try_into().expect("pg FLOAT8 is always 8 bytes"));
            rmp::encode::write_f64(out, v).expect("write to Vec never fails");
        }
        oid::TEXT | oid::VARCHAR => {
            out.extend_from_slice(&pg_buf);
        }
        oid::TIMESTAMP | oid::TIMESTAMPTZ => {
            encode_datetime(&pg_buf, out);
        }
        oid::UUID => {
            rmp::encode::write_ext_meta(out, 16, 0x02).expect("write to Vec never fails");
            out.extend_from_slice(&pg_buf[..16]);
        }
        _ => return false,
    }
    true
}

/// Maps an OID to a `pg::types::Type`. Returns `None` for unsupported OIDs.
fn pg_type_from_oid(type_oid: u32) -> Option<pg::types::Type> {
    Some(match type_oid {
        oid::BOOL => pg::types::Type::BOOL,
        oid::INT2 => pg::types::Type::INT2,
        oid::INT4 => pg::types::Type::INT4,
        oid::INT8 => pg::types::Type::INT8,
        oid::FLOAT4 => pg::types::Type::FLOAT4,
        oid::FLOAT8 => pg::types::Type::FLOAT8,
        oid::TEXT => pg::types::Type::TEXT,
        oid::VARCHAR => pg::types::Type::VARCHAR,
        oid::TIMESTAMP => pg::types::Type::TIMESTAMP,
        oid::TIMESTAMPTZ => pg::types::Type::TIMESTAMPTZ,
        oid::UUID => pg::types::Type::UUID,
        _ => return None,
    })
}

/// Microseconds between Unix epoch (1970) and PG epoch (2000).
const PG_EPOCH_DIFF_US: i64 = 946_684_800_000_000;

/// Encodes a PG TIMESTAMP/TIMESTAMPTZ as Tarantool msgpack Datetime (`MP_EXT` type 0x04).
fn encode_datetime(pg_buf: &[u8], out: &mut Vec<u8>) {
    let pg_us = i64::from_be_bytes(pg_buf[..8].try_into().expect("pg INT8 is always 8 bytes"));
    let unix_us = pg_us + PG_EPOCH_DIFF_US;
    let seconds = unix_us.div_euclid(1_000_000);
    let nsec = i32::try_from(unix_us.rem_euclid(1_000_000) * 1_000)
        .expect("nanoseconds should fit in i32");

    if nsec == 0 {
        rmp::encode::write_ext_meta(out, 8, 0x04).expect("write to Vec never fails");
        out.extend_from_slice(&seconds.to_le_bytes());
    } else {
        rmp::encode::write_ext_meta(out, 16, 0x04).expect("write to Vec never fails");
        out.extend_from_slice(&seconds.to_le_bytes());
        out.extend_from_slice(&nsec.to_le_bytes());
        out.extend_from_slice(&0i16.to_le_bytes());
        out.extend_from_slice(&0i16.to_le_bytes());
    }
}

/// Computes `bucket_id` from query parameters using distribution key metadata.
///
/// Encodes each distribution key parameter as msgpack, hashes with `MurmurHash3` (seed 13),
/// and maps to a bucket in range `[1, bucket_count]`.
///
/// Returns `None` if any parameter has an unsupported OID.
pub(crate) fn compute_bucket_id(
    params: &[&(dyn pg::types::ToSql + Sync)],
    dk_meta: &DkMeta,
    bucket_count: u64,
) -> Option<u64> {
    let mut buf = Vec::with_capacity(dk_meta.keys.len() * 16);
    for &(idx, type_oid) in &dk_meta.keys {
        // idx is 1-based ($1=1, $2=2); params is 0-based.
        if !encode_param_msgpack(params[(idx - 1) as usize], type_oid, &mut buf) {
            return None;
        }
    }
    let hash = murmur3::murmur3_32(&mut std::io::Cursor::new(&buf), 13)
        .expect("hashing in-memory buffer should not fail");
    let bucket_id = u64::from(hash) % bucket_count + 1;
    Some(bucket_id)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::metadata::DkMeta;

    /// Regression: `dk_meta.keys` stores 1-based PostgreSQL parameter indices
    /// (`$1` → idx=1, `$2` → idx=2), but `params` is 0-based.  Before the fix,
    /// `params[idx as usize]` accessed `params[1]` for `$1`, causing an
    /// index-out-of-bounds panic on the common single-parameter case.
    #[test]
    fn compute_bucket_id_maps_1based_key_to_0based_params() {
        let dk = DkMeta {
            tier: Some("default".to_string()),
            // $1 → 1-based idx == 1; params has exactly one element.
            // Before the fix: params[1] panics (len == 1, index == 1).
            keys: vec![(1, oid::INT8)],
        };
        let v: i64 = 42;
        let params: Vec<&(dyn pg::types::ToSql + Sync)> = vec![&v];
        assert!(
            compute_bucket_id(&params, &dk, 3000).is_some(),
            "must not panic: 1-based idx 1 must map to params[0], not params[1]"
        );
    }
}