Struct google_cloud_spanner::session::SessionHandle

pub struct SessionHandle {
    pub session: Session,
    pub spanner_client: Client,
    /* private fields */
}

Session

Fields

session: Session

spanner_client: Client

Implementations

impl SessionHandle

pub async fn invalidate_if_needed<T>(
    &mut self,
    arg: Result<T, Status>
) -> Result<T, Status>

Examples found in repository

src/reader.rs (line 50)

    async fn read(
        &self,
        session: &mut SessionHandle,
        option: Option<CallOptions>,
    ) -> Result<Response<Streaming<PartialResultSet>>, Status> {
        let option = option.unwrap_or_default();
        let client = &mut session.spanner_client;
        let result = client
            .execute_streaming_sql(self.request.clone(), option.cancel, option.retry)
            .await;
        return session.invalidate_if_needed(result).await;
    }

    fn update_token(&mut self, resume_token: Vec<u8>) {
        self.request.resume_token = resume_token;
    }

    fn can_retry(&self) -> bool {
        !self.request.resume_token.is_empty()
    }
}

pub struct TableReader {
    pub request: ReadRequest,
}

#[async_trait]
impl Reader for TableReader {
    async fn read(
        &self,
        session: &mut SessionHandle,
        option: Option<CallOptions>,
    ) -> Result<Response<Streaming<PartialResultSet>>, Status> {
        let option = option.unwrap_or_default();
        let client = &mut session.spanner_client;
        let result = client
            .streaming_read(self.request.clone(), option.cancel, option.retry)
            .await;
        return session.invalidate_if_needed(result).await;
    }

src/transaction_rw.rs (line 140)

    async fn begin_internal(
        mut session: ManagedSession,
        mode: transaction_options::Mode,
        options: CallOptions,
    ) -> Result<ReadWriteTransaction, BeginError> {
        let request = BeginTransactionRequest {
            session: session.session.name.to_string(),
            options: Some(TransactionOptions { mode: Some(mode) }),
            request_options: Transaction::create_request_options(options.priority),
        };
        let result = session
            .spanner_client
            .begin_transaction(request, options.cancel, options.retry)
            .await;
        let response = match session.invalidate_if_needed(result).await {
            Ok(response) => response,
            Err(err) => {
                return Err(BeginError { status: err, session });
            }
        };
        let tx = response.into_inner();
        Ok(ReadWriteTransaction {
            base_tx: Transaction {
                session: Some(session),
                sequence_number: AtomicI64::new(0),
                transaction_selector: TransactionSelector {
                    selector: Some(transaction_selector::Selector::Id(tx.id.clone())),
                },
            },
            tx_id: tx.id,
            wb: vec![],
        })
    }

    pub fn buffer_write(&mut self, ms: Vec<Mutation>) {
        self.wb.extend_from_slice(&ms)
    }

    pub async fn update(&mut self, stmt: Statement) -> Result<i64, Status> {
        self.update_with_option(stmt, QueryOptions::default()).await
    }

    pub async fn update_with_option(&mut self, stmt: Statement, options: QueryOptions) -> Result<i64, Status> {
        let request = ExecuteSqlRequest {
            session: self.get_session_name(),
            transaction: Some(self.transaction_selector.clone()),
            sql: stmt.sql.to_string(),
            params: Some(prost_types::Struct { fields: stmt.params }),
            param_types: stmt.param_types,
            resume_token: vec![],
            query_mode: options.mode.into(),
            partition_token: vec![],
            seqno: self.sequence_number.fetch_add(1, Ordering::Relaxed),
            query_options: options.optimizer_options,
            request_options: Transaction::create_request_options(options.call_options.priority),
        };

        let session = self.as_mut_session();
        let result = session
            .spanner_client
            .execute_sql(request, options.call_options.cancel, options.call_options.retry)
            .await;
        let response = session.invalidate_if_needed(result).await?;
        Ok(extract_row_count(response.into_inner().stats))
    }

    pub async fn batch_update(&mut self, stmt: Vec<Statement>) -> Result<Vec<i64>, Status> {
        self.batch_update_with_option(stmt, QueryOptions::default()).await
    }

    pub async fn batch_update_with_option(
        &mut self,
        stmt: Vec<Statement>,
        options: QueryOptions,
    ) -> Result<Vec<i64>, Status> {
        let request = ExecuteBatchDmlRequest {
            session: self.get_session_name(),
            transaction: Some(self.transaction_selector.clone()),
            seqno: self.sequence_number.fetch_add(1, Ordering::Relaxed),
            request_options: Transaction::create_request_options(options.call_options.priority),
            statements: stmt
                .into_iter()
                .map(|x| execute_batch_dml_request::Statement {
                    sql: x.sql,
                    params: Some(Struct { fields: x.params }),
                    param_types: x.param_types,
                })
                .collect(),
        };

        let session = self.as_mut_session();
        let result = session
            .spanner_client
            .execute_batch_dml(request, options.call_options.cancel, options.call_options.retry)
            .await;
        let response = session.invalidate_if_needed(result).await?;
        Ok(response
            .into_inner()
            .result_sets
            .into_iter()
            .map(|x| extract_row_count(x.stats))
            .collect())
    }

    pub async fn end<S, E>(
        &mut self,
        result: Result<S, E>,
        options: Option<CommitOptions>,
    ) -> Result<(Option<Timestamp>, S), E>
    where
        E: TryAs<Status> + From<Status>,
    {
        let opt = options.unwrap_or_default();
        match result {
            Ok(success) => {
                let cr = self.commit(opt).await?;
                Ok((cr.commit_timestamp.map(|e| e.into()), success))
            }
            Err(err) => {
                if let Some(status) = err.try_as() {
                    // can't rollback. should retry
                    if status.code() == Code::Aborted {
                        return Err(err);
                    }
                }
                let _ = self.rollback(opt.call_options.cancel, opt.call_options.retry).await;
                Err(err)
            }
        }
    }

    pub(crate) async fn finish<T, E>(
        &mut self,
        result: Result<T, E>,
        options: Option<CommitOptions>,
    ) -> Result<(Option<Timestamp>, T), (E, Option<ManagedSession>)>
    where
        E: TryAs<Status> + From<Status>,
    {
        let opt = options.unwrap_or_default();

        return match result {
            Ok(s) => match self.commit(opt).await {
                Ok(c) => Ok((c.commit_timestamp.map(|ts| ts.into()), s)),
                // Retry the transaction using the same session on ABORT error.
                // Cloud Spanner will create the new transaction with the previous
                // one's wound-wait priority.
                Err(e) => Err((E::from(e), self.take_session())),
            },

            // Rollback the transaction unless the error occurred during the
            // commit. Executing a rollback after a commit has failed will
            // otherwise cause an error. Note that transient errors, such as
            // UNAVAILABLE, are already handled in the gRPC layer and do not show
            // up here. Context errors (deadline exceeded / canceled) during
            // commits are also not rolled back.
            Err(err) => {
                let status = match err.try_as() {
                    Some(status) => status,
                    None => {
                        let _ = self.rollback(opt.call_options.cancel, opt.call_options.retry).await;
                        return Err((err, self.take_session()));
                    }
                };
                match status.code() {
                    Code::Aborted => Err((err, self.take_session())),
                    _ => {
                        let _ = self.rollback(opt.call_options.cancel, opt.call_options.retry).await;
                        return Err((err, self.take_session()));
                    }
                }
            }
        };
    }

    pub(crate) async fn commit(&mut self, options: CommitOptions) -> Result<CommitResponse, Status> {
        let tx_id = self.tx_id.clone();
        let mutations = self.wb.to_vec();
        let session = self.as_mut_session();
        commit(session, mutations, TransactionId(tx_id), options).await
    }

    pub(crate) async fn rollback(
        &mut self,
        cancel: Option<CancellationToken>,
        retry: Option<RetrySetting>,
    ) -> Result<(), Status> {
        let request = RollbackRequest {
            transaction_id: self.tx_id.clone(),
            session: self.get_session_name(),
        };
        let session = self.as_mut_session();
        let result = session.spanner_client.rollback(request, cancel, retry).await;
        session.invalidate_if_needed(result).await?.into_inner();
        Ok(())
    }
}

pub(crate) async fn commit(
    session: &mut ManagedSession,
    ms: Vec<Mutation>,
    tx: commit_request::Transaction,
    commit_options: CommitOptions,
) -> Result<CommitResponse, Status> {
    let request = CommitRequest {
        session: session.session.name.to_string(),
        mutations: ms,
        transaction: Some(tx),
        request_options: Transaction::create_request_options(commit_options.call_options.priority),
        return_commit_stats: commit_options.return_commit_stats,
    };
    let result = session
        .spanner_client
        .commit(request, commit_options.call_options.cancel, commit_options.call_options.retry)
        .await;
    let response = session.invalidate_if_needed(result).await;
    match response {
        Ok(r) => Ok(r.into_inner()),
        Err(s) => Err(s),
    }
}

src/transaction_ro.rs (line 86)

    pub async fn begin(
        mut session: ManagedSession,
        tb: TimestampBound,
        options: CallOptions,
    ) -> Result<ReadOnlyTransaction, Status> {
        let request = BeginTransactionRequest {
            session: session.session.name.to_string(),
            options: Some(TransactionOptions {
                mode: Some(transaction_options::Mode::ReadOnly(tb.into())),
            }),
            request_options: Transaction::create_request_options(options.priority),
        };

        let result = session
            .spanner_client
            .begin_transaction(request, options.cancel, options.retry)
            .await;
        match session.invalidate_if_needed(result).await {
            Ok(response) => {
                let tx = response.into_inner();
                let rts = tx.read_timestamp.unwrap();
                let st: SystemTime = rts.try_into().unwrap();
                Ok(ReadOnlyTransaction {
                    base_tx: Transaction {
                        session: Some(session),
                        sequence_number: AtomicI64::new(0),
                        transaction_selector: TransactionSelector {
                            selector: Some(transaction_selector::Selector::Id(tx.id)),
                        },
                    },
                    rts: Some(OffsetDateTime::from(st)),
                })
            }
            Err(e) => Err(e),
        }
    }
}

pub struct Partition<T: Reader> {
    pub reader: T,
}

/// BatchReadOnlyTransaction is a ReadOnlyTransaction that allows for exporting
/// arbitrarily large amounts of data from Cloud Spanner databases.
/// BatchReadOnlyTransaction partitions a read/query request. Read/query request
/// can then be executed independently over each partition while observing the
/// same snapshot of the database.
pub struct BatchReadOnlyTransaction {
    base_tx: ReadOnlyTransaction,
}

impl Deref for BatchReadOnlyTransaction {
    type Target = ReadOnlyTransaction;

    fn deref(&self) -> &Self::Target {
        &self.base_tx
    }
}

impl DerefMut for BatchReadOnlyTransaction {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.base_tx
    }
}

impl BatchReadOnlyTransaction {
    pub async fn begin(
        session: ManagedSession,
        tb: TimestampBound,
        options: CallOptions,
    ) -> Result<BatchReadOnlyTransaction, Status> {
        let tx = ReadOnlyTransaction::begin(session, tb, options).await?;
        Ok(BatchReadOnlyTransaction { base_tx: tx })
    }

    /// partition_read returns a list of Partitions that can be used to read rows from
    /// the database. These partitions can be executed across multiple processes,
    /// even across different machines. The partition size and count hints can be
    /// configured using PartitionOptions.
    pub async fn partition_read(
        &mut self,
        table: &str,
        columns: &[&str],
        keys: impl Into<KeySet> + Clone,
    ) -> Result<Vec<Partition<TableReader>>, Status> {
        self.partition_read_with_option(table, columns, keys, None, ReadOptions::default())
            .await
    }

    /// partition_read returns a list of Partitions that can be used to read rows from
    /// the database. These partitions can be executed across multiple processes,
    /// even across different machines. The partition size and count hints can be
    /// configured using PartitionOptions.
    pub async fn partition_read_with_option(
        &mut self,
        table: &str,
        columns: &[&str],
        keys: impl Into<KeySet> + Clone,
        po: Option<PartitionOptions>,
        ro: ReadOptions,
    ) -> Result<Vec<Partition<TableReader>>, Status> {
        let columns: Vec<String> = columns.iter().map(|x| x.to_string()).collect();
        let inner_keyset = keys.into().inner;
        let request = PartitionReadRequest {
            session: self.get_session_name(),
            transaction: Some(self.transaction_selector.clone()),
            table: table.to_string(),
            index: ro.index.clone(),
            columns: columns.clone(),
            key_set: Some(inner_keyset.clone()),
            partition_options: po,
        };
        let result = match self
            .as_mut_session()
            .spanner_client
            .partition_read(request, ro.call_options.cancel, ro.call_options.retry)
            .await
        {
            Ok(r) => Ok(r
                .into_inner()
                .partitions
                .into_iter()
                .map(|x| Partition {
                    reader: TableReader {
                        request: ReadRequest {
                            session: self.get_session_name(),
                            transaction: Some(self.transaction_selector.clone()),
                            table: table.to_string(),
                            index: ro.index.clone(),
                            columns: columns.clone(),
                            key_set: Some(inner_keyset.clone()),
                            limit: ro.limit,
                            resume_token: vec![],
                            partition_token: x.partition_token,
                            request_options: Transaction::create_request_options(ro.call_options.priority),
                        },
                    },
                })
                .collect()),
            Err(e) => Err(e),
        };
        self.as_mut_session().invalidate_if_needed(result).await
    }

    /// partition_query returns a list of Partitions that can be used to execute a query against the database.
    pub async fn partition_query(&mut self, stmt: Statement) -> Result<Vec<Partition<StatementReader>>, Status> {
        self.partition_query_with_option(stmt, None, QueryOptions::default())
            .await
    }

    /// partition_query returns a list of Partitions that can be used to execute a query against the database.
    pub async fn partition_query_with_option(
        &mut self,
        stmt: Statement,
        po: Option<PartitionOptions>,
        qo: QueryOptions,
    ) -> Result<Vec<Partition<StatementReader>>, Status> {
        let request = PartitionQueryRequest {
            session: self.get_session_name(),
            transaction: Some(self.transaction_selector.clone()),
            sql: stmt.sql.clone(),
            params: Some(prost_types::Struct {
                fields: stmt.params.clone(),
            }),
            param_types: stmt.param_types.clone(),
            partition_options: po,
        };
        let result = match self
            .as_mut_session()
            .spanner_client
            .partition_query(request.clone(), qo.call_options.cancel.clone(), qo.call_options.retry.clone())
            .await
        {
            Ok(r) => Ok(r
                .into_inner()
                .partitions
                .into_iter()
                .map(|x| Partition {
                    reader: StatementReader {
                        request: ExecuteSqlRequest {
                            session: self.get_session_name(),
                            transaction: Some(self.transaction_selector.clone()),
                            sql: stmt.sql.clone(),
                            params: Some(prost_types::Struct {
                                fields: stmt.params.clone(),
                            }),
                            param_types: stmt.param_types.clone(),
                            resume_token: vec![],
                            query_mode: 0,
                            partition_token: x.partition_token,
                            seqno: 0,
                            query_options: qo.optimizer_options.clone(),
                            request_options: Transaction::create_request_options(qo.call_options.priority),
                        },
                    },
                })
                .collect()),
            Err(e) => Err(e),
        };
        self.as_mut_session().invalidate_if_needed(result).await
    }