laminar_connectors/ibmmq/sink/
mod.rs1mod message;
10
11use std::time::Duration;
12
13use arrow_array::RecordBatch;
14use arrow_schema::SchemaRef;
15use async_trait::async_trait;
16use mqi::prelude::*;
17use mqi::types::{MessageFormat, MQCMHO, MQPMO, MQSMPO};
18use mqi::{constants, structs, Properties};
19
20use super::config::{IbmMqSinkConfig, Persistence};
21use super::library;
22use super::metrics::SinkMetrics;
23use super::session::{classify, Session};
24use super::worker::Worker;
25use crate::config::ConnectorConfig;
26use crate::connector::{
27 ConnectorTaskOwner, ConnectorTaskTracker, SinkConnector, SinkConsistency, SinkContract,
28 SinkInputMode, SinkTopology, WriteResult,
29};
30use crate::error::ConnectorError;
31use crate::serde::RecordSerializer;
32use message::{base_descriptor, payload_schema, Outgoing};
33
34const CLOSE_TIMEOUT: Duration = Duration::from_secs(30);
35
36pub struct IbmMqSink {
38 schema: SchemaRef,
39 registry: Option<prometheus::Registry>,
40 config: Option<IbmMqSinkConfig>,
41 writer: Option<Writer>,
42 worker: Option<Worker<Session>>,
43 metrics: Option<SinkMetrics>,
44 failure: Option<String>,
45 task_owner: ConnectorTaskOwner,
46 task_tracker: ConnectorTaskTracker,
47}
48
49struct Writer {
50 serializer: Box<dyn RecordSerializer>,
51 payload_indices: Vec<usize>,
52 base: libmqm_sys::MQMD,
53}
54
55struct PutFailure {
57 committed: usize,
58 error: ConnectorError,
59}
60
61fn put_all(session: &Session, messages: &[Outgoing], unit: usize) -> Result<usize, PutFailure> {
63 let mut committed = 0;
64 for part in messages.chunks(unit) {
65 for outgoing in part {
66 if let Err(error) = put_one(session, outgoing) {
67 let backed_out = session.backout();
68 let error = match backed_out {
69 Ok(()) => error,
70 Err(backout) => ConnectorError::outcome_unknown(
71 format!("{error}; MQBACK also failed: {backout}"),
72 true,
73 ),
74 };
75 return Err(PutFailure { committed, error });
76 }
77 }
78 session
79 .commit()
80 .map_err(|error| PutFailure { committed, error })?;
81 committed += part.len();
82 }
83 Ok(committed)
84}
85
86fn put_one(session: &Session, outgoing: &Outgoing) -> Result<(), ConnectorError> {
87 let handle = Properties::new(session.connection(), MQCMHO::default())
88 .map_err(|error| classify("MQCRTMH", &error))?;
89 for (name, value) in &outgoing.properties {
90 handle
91 .set_property(name.as_str(), value, MQSMPO::default())
92 .warn_as_error()
93 .map_err(|error| classify(&format!("MQSETMP '{name}'"), &error))?;
94 }
95 let mut options: MQPMO = constants::MQPMO_SYNCPOINT | constants::MQPMO_FAIL_IF_QUIESCING;
96 if outgoing.new_message_id {
97 options |= constants::MQPMO_NEW_MSG_ID;
98 }
99 let md = structs::MQMD::new(outgoing.md);
100 let format = MessageFormat::from_mqmd(&md);
101 let mut handle = handle;
102 session
103 .queue
104 .put_message(
105 &(options, md, &mut handle),
106 &(outgoing.data.as_slice(), format),
107 )
108 .warn_as_error()
109 .map_err(|error| match error.2 {
110 constants::MQRC_MSG_TOO_BIG_FOR_Q | constants::MQRC_MSG_TOO_BIG_FOR_Q_MGR => {
111 ConnectorError::ConfigurationError(format!(
112 "MQPUT: a {}-byte message exceeds the queue's MAXMSGL: {error}",
113 outgoing.data.len()
114 ))
115 }
116 _ => classify("MQPUT", &error),
117 })
118}
119
120impl IbmMqSink {
121 #[must_use]
123 pub fn new(schema: SchemaRef, registry: Option<&prometheus::Registry>) -> Self {
124 let (task_owner, task_tracker) = ConnectorTaskOwner::new();
125 Self {
126 schema,
127 registry: registry.cloned(),
128 config: None,
129 writer: None,
130 worker: None,
131 metrics: None,
132 failure: None,
133 task_owner,
134 task_tracker,
135 }
136 }
137
138 fn writer(parsed: &IbmMqSinkConfig, input: &SchemaRef) -> Result<Writer, ConnectorError> {
139 let payload = payload_schema(parsed, input)?;
140 crate::serde::schema_contract::validate_writer(parsed.format, &payload)?;
141 let payload_indices = payload
142 .fields()
143 .iter()
144 .map(|field| input.index_of(field.name()))
145 .collect::<Result<_, _>>()
146 .map_err(|error| ConnectorError::SchemaMismatch(error.to_string()))?;
147 Ok(Writer {
148 serializer: crate::serde::create_serializer(parsed.format)?,
149 base: base_descriptor(parsed, &payload)?,
150 payload_indices,
151 })
152 }
153
154 fn fail(&mut self, error: ConnectorError) -> ConnectorError {
155 tracing::error!(%error, "IBM MQ sink generation ended");
156 self.worker = None;
157 self.failure.get_or_insert(error.to_string());
158 error
159 }
160
161 async fn put(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
162 if let Some(cause) = &self.failure {
163 return Err(ConnectorError::InvalidState {
164 expected: "a live IBM MQ sink".into(),
165 actual: cause.clone(),
166 });
167 }
168 let (Some(config), Some(writer), Some(worker), Some(metrics)) =
169 (&self.config, &self.writer, &self.worker, &self.metrics)
170 else {
171 return Err(ConnectorError::InvalidState {
172 expected: "an open IBM MQ sink".into(),
173 actual: "open() has not completed".into(),
174 });
175 };
176 let payload = batch
177 .project(&writer.payload_indices)
178 .map_err(|error| ConnectorError::SchemaMismatch(error.to_string()))?;
179 let records = writer.serializer.serialize(&payload)?;
180 if records.len() != batch.num_rows() {
181 return Err(ConnectorError::Internal(format!(
182 "IBM MQ serializer returned {} records for {} rows",
183 records.len(),
184 batch.num_rows()
185 )));
186 }
187 let outgoing = message::messages(config, &writer.base, batch, records)?;
188 let bytes: u64 = outgoing
189 .iter()
190 .map(|message| message.data.len() as u64)
191 .sum();
192 let count = outgoing.len();
193 let unit = config.max_uow_messages;
194 let result = worker
195 .call(
196 "MQPUT/MQCMIT",
197 config.connection.operation_timeout,
198 move |session| Ok(put_all(session, &outgoing, unit)),
199 )
200 .await?;
201 match result {
202 Ok(committed) => {
203 metrics.put.inc_by(committed as u64);
204 metrics.bytes.inc_by(bytes);
205 metrics.commits.inc_by(committed.div_ceil(unit) as u64);
206 Ok(WriteResult::new(count, bytes))
207 }
208 Err(PutFailure { committed, error }) => {
209 metrics.backouts.inc();
210 if committed > 0 || error.is_outcome_unknown() {
211 metrics.ambiguous.inc();
212 Err(ConnectorError::outcome_unknown(
213 format!(
214 "IBM MQ batch partially committed ({committed} of {count}): {error}"
215 ),
216 true,
217 ))
218 } else {
219 Err(error)
220 }
221 }
222 }
223 }
224}
225
226#[async_trait]
227impl SinkConnector for IbmMqSink {
228 fn terminal_task_tracker(&self) -> Option<ConnectorTaskTracker> {
229 Some(self.task_tracker.clone())
230 }
231
232 fn contract(&self, config: &ConnectorConfig) -> Result<SinkContract, ConnectorError> {
233 let parsed = IbmMqSinkConfig::from_config(config)?;
234 let consistency = match parsed.persistence {
236 Persistence::Persistent => SinkConsistency::DurableAtLeastOnce,
237 Persistence::NotPersistent | Persistence::AsQueueDefinition => {
238 SinkConsistency::Ephemeral
239 }
240 };
241 Ok(SinkContract::new(
242 consistency,
243 SinkTopology::MultiWriter,
244 SinkInputMode::AppendOnly,
245 ))
246 }
247
248 async fn resolve_schema(
249 &mut self,
250 config: &ConnectorConfig,
251 input: SchemaRef,
252 ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
253 use crate::schema::resolution::{logical_binding, SchemaDirection, SchemaOrigin};
254 let parsed = IbmMqSinkConfig::from_config(config)?;
255 Self::writer(&parsed, &input)?;
256 logical_binding(config, SchemaDirection::Sink, SchemaOrigin::Query, &input)
257 }
258
259 async fn open(&mut self, config: &ConnectorConfig) -> Result<(), ConnectorError> {
260 let parsed = IbmMqSinkConfig::from_config(config)?;
261 let input = config
262 .arrow_schema()
263 .or_else(|| {
264 config
265 .schema_binding()
266 .map(|binding| std::sync::Arc::new(binding.logical.clone()))
267 })
268 .unwrap_or_else(|| self.schema.clone());
269 let writer = Self::writer(&parsed, &input)?;
270 let library = library::load(parsed.connection.client_library.as_deref())?;
271 let metrics = SinkMetrics::new(&parsed.queue, self.registry.as_ref())?;
272 let guard = self.task_owner.track().ok_or_else(|| {
273 ConnectorError::Internal("IBM MQ sink task generation is retired".into())
274 })?;
275 let session_config = parsed.clone();
276 let worker = Worker::start(
277 &format!("ibmmq-sink-{}", parsed.queue),
278 guard,
279 parsed.connection.operation_timeout,
280 move || {
281 Session::open(
282 &session_config.connection,
283 library,
284 &session_config.queue,
285 session_config.target_queue_manager.as_deref(),
286 constants::MQOO_OUTPUT,
287 )
288 },
289 )
290 .await?;
291 tracing::info!(
292 queue = %parsed.queue,
293 persistence = ?parsed.persistence,
294 connection = ?parsed.connection,
295 "IBM MQ sink putting under syncpoint"
296 );
297 self.schema = input;
298 self.writer = Some(writer);
299 self.worker = Some(worker);
300 self.metrics = Some(metrics);
301 self.config = Some(parsed);
302 self.failure = None;
303 Ok(())
304 }
305
306 async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
307 match self.put(batch).await {
308 Ok(result) => Ok(result),
309 Err(error) if self.failure.is_some() => Err(error),
310 Err(error) => Err(self.fail(error)),
311 }
312 }
313
314 fn schema(&self) -> SchemaRef {
315 self.schema.clone()
316 }
317
318 fn suggested_write_timeout(&self) -> Duration {
319 self.config
320 .as_ref()
321 .map_or(Duration::from_secs(30), |config| {
322 config.connection.operation_timeout
323 })
324 }
325
326 async fn close(&mut self) -> Result<(), ConnectorError> {
327 match self.worker.take() {
328 Some(worker) => worker.stop(CLOSE_TIMEOUT).await,
329 None => Ok(()),
330 }
331 }
332}
333
334#[cfg(test)]
335mod tests;