use std::sync::Arc;
use ydb_grpc::ydb_proto::topic::TransactionIdentity;
use tracing::instrument;
use crate::YdbResult;
use crate::client_query::Transaction;
use crate::client_query::hooks::{QueryTxCommitStatus, QueryTxHook};
use crate::client_topic::compression::Executor;
use crate::client_topic::topicwriter::message::TopicWriterMessage;
use crate::client_topic::topicwriter::writer::TopicWriter;
use crate::grpc_connection_manager::GrpcConnectionManager;
use super::writer_tx_options::TopicWriterTxOptions;
pub struct TopicWriterTx {
inner: Arc<TopicWriter>,
}
struct WriterTxHook {
writer: Arc<TopicWriter>,
}
impl WriterTxHook {
#[instrument(name = "ydb.TopicWriterTx.Flush", skip_all, fields(db.system.name = "ydb"), err)]
async fn flush(&self) -> YdbResult<()> {
self.writer.flush_inner().await
}
}
#[async_trait::async_trait]
impl QueryTxHook for WriterTxHook {
async fn before_commit(&mut self) -> YdbResult<()> {
self.flush().await
}
fn after_commit(&mut self, _status: QueryTxCommitStatus) {}
}
impl TopicWriterTx {
pub(crate) async fn new(
options: TopicWriterTxOptions,
connection_manager: GrpcConnectionManager,
executor: Arc<dyn Executor>,
tx: &mut Transaction,
) -> YdbResult<Self> {
let (session_id, transaction_id) = tx.identity().await?;
let tx_identity = TransactionIdentity {
id: transaction_id,
session: session_id,
};
let options = options.into_non_tx_options();
let writer =
TopicWriter::with_tx_identity(options, connection_manager, executor, tx_identity)
.await?;
let inner = Arc::new(writer);
tx.register_hook(WriterTxHook {
writer: inner.clone(),
});
Ok(Self { inner })
}
#[instrument(name = "ydb.TopicWriterTx.Write", skip_all, fields(db.system.name = "ydb"), err)]
pub async fn write(&mut self, message: TopicWriterMessage) -> YdbResult<()> {
self.inner.write_inner(message).await
}
}