sqlx_scylladb_core/
statement.rs

1use std::{borrow::Cow, sync::Arc};
2
3use scylla::{response::query_result::ColumnSpecs, statement::prepared::PreparedStatement};
4use sqlx::{ColumnIndex, Error, Statement};
5use sqlx_core::{HashMap, ext::ustr::UStr, impl_statement_query};
6
7use crate::{ScyllaDB, ScyllaDBArguments, ScyllaDBColumn, ScyllaDBError, ScyllaDBTypeInfo};
8
9/// Implementation of [sqlx::Statement] for ScyllaDB.
10#[derive(Clone)]
11pub struct ScyllaDBStatement<'q> {
12    pub(crate) sql: Cow<'q, str>,
13    pub(crate) prepared_statement: PreparedStatement,
14    pub(crate) metadata: ScyllaDBStatementMetadata,
15    pub(crate) is_affect_statement: bool,
16}
17
18impl<'q> Statement<'q> for ScyllaDBStatement<'q> {
19    type Database = ScyllaDB;
20
21    fn to_owned(&self) -> ScyllaDBStatement<'static> {
22        ScyllaDBStatement::<'static> {
23            sql: Cow::Owned(self.sql.clone().into_owned()),
24            prepared_statement: self.prepared_statement.clone(),
25            metadata: self.metadata.clone(),
26            is_affect_statement: self.is_affect_statement,
27        }
28    }
29
30    fn sql(&self) -> &str {
31        &self.sql
32    }
33
34    fn parameters(&self) -> Option<sqlx::Either<&[ScyllaDBTypeInfo], usize>> {
35        Some(sqlx::Either::Right(self.metadata.parameters))
36    }
37
38    fn columns(&self) -> &[ScyllaDBColumn] {
39        &self.metadata.columns
40    }
41
42    impl_statement_query!(ScyllaDBArguments);
43}
44
45impl ColumnIndex<ScyllaDBStatement<'_>> for &'_ str {
46    fn index(&self, statement: &ScyllaDBStatement<'_>) -> Result<usize, Error> {
47        statement
48            .metadata
49            .column_names
50            .get(*self)
51            .ok_or_else(|| Error::ColumnNotFound((*self).into()))
52            .copied()
53    }
54}
55
56#[derive(Debug, Default, Clone)]
57pub(crate) struct ScyllaDBStatementMetadata {
58    pub(crate) columns: Arc<Vec<ScyllaDBColumn>>,
59    pub(crate) column_names: Arc<HashMap<UStr, usize>>,
60    pub(crate) parameters: usize,
61}
62
63impl ScyllaDBStatementMetadata {
64    pub(crate) fn from_column_specs(column_specs: ColumnSpecs) -> Result<Self, ScyllaDBError> {
65        let parameters = column_specs.len();
66        let mut columns = Vec::with_capacity(parameters);
67        let mut column_names = HashMap::with_capacity(parameters);
68        for (i, column_spec) in column_specs.iter().enumerate() {
69            let name = UStr::new(column_spec.name());
70            let column_type = column_spec.typ();
71            let type_info = ScyllaDBTypeInfo::from_column_type(column_type)?;
72
73            column_names.insert(name.clone(), i);
74            columns.push(ScyllaDBColumn {
75                ordinal: i,
76                name,
77                type_info,
78                column_type: column_type.clone().into_owned(),
79            })
80        }
81
82        let columns = Arc::new(columns);
83        let column_names = Arc::new(column_names);
84
85        let metadata = ScyllaDBStatementMetadata {
86            columns,
87            column_names,
88            parameters,
89        };
90
91        Ok(metadata)
92    }
93}