1use std::collections::BTreeMap;
16use std::sync::Arc;
17
18use ahash::{HashMap, HashMapExt};
19use api::greptime_proto;
20use api::v1::{ColumnDataType, ColumnSchema, RowInsertRequest, Rows, SemanticType};
21use common_time::timestamp::TimeUnit;
22use pipeline::{
23 identity_pipeline, unwrap_or_continue_if_err, ContextReq, DispatchedTo, Pipeline,
24 PipelineContext, PipelineDefinition, PipelineExecOutput, SchemaInfo, TransformedOutput,
25 TransformerMode, GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME,
26};
27use session::context::{Channel, QueryContextRef};
28use snafu::ResultExt;
29use vrl::value::Value as VrlValue;
30
31use crate::error::{CatalogSnafu, PipelineSnafu, Result};
32use crate::http::event::PipelineIngestRequest;
33use crate::metrics::{
34 METRIC_FAILURE_VALUE, METRIC_HTTP_LOGS_TRANSFORM_ELAPSED, METRIC_SUCCESS_VALUE,
35};
36use crate::query_handler::PipelineHandlerRef;
37
38macro_rules! push_to_map {
39 ($map:expr, $key:expr, $value:expr, $capacity:expr) => {
40 $map.entry($key)
41 .or_insert_with(|| Vec::with_capacity($capacity))
42 .push($value);
43 };
44}
45
46pub async fn get_pipeline(
48 pipeline_def: &PipelineDefinition,
49 handler: &PipelineHandlerRef,
50 query_ctx: &QueryContextRef,
51) -> Result<Arc<Pipeline>> {
52 match pipeline_def {
53 PipelineDefinition::Resolved(pipeline) => Ok(pipeline.clone()),
54 PipelineDefinition::ByNameAndValue((name, version)) => {
55 handler
56 .get_pipeline(name, *version, query_ctx.clone())
57 .await
58 }
59 _ => {
60 unreachable!("Never call get_pipeline on identity.")
61 }
62 }
63}
64
65pub(crate) async fn run_pipeline(
66 handler: &PipelineHandlerRef,
67 pipeline_ctx: &PipelineContext<'_>,
68 pipeline_req: PipelineIngestRequest,
69 query_ctx: &QueryContextRef,
70 is_top_level: bool,
71) -> Result<ContextReq> {
72 if pipeline_ctx.pipeline_definition.is_identity() {
73 run_identity_pipeline(handler, pipeline_ctx, pipeline_req, query_ctx).await
74 } else {
75 run_custom_pipeline(handler, pipeline_ctx, pipeline_req, query_ctx, is_top_level).await
76 }
77}
78
79async fn run_identity_pipeline(
80 handler: &PipelineHandlerRef,
81 pipeline_ctx: &PipelineContext<'_>,
82 pipeline_req: PipelineIngestRequest,
83 query_ctx: &QueryContextRef,
84) -> Result<ContextReq> {
85 let PipelineIngestRequest {
86 table: table_name,
87 values: data_array,
88 } = pipeline_req;
89 let table = if pipeline_ctx.channel == Channel::Prometheus {
90 None
91 } else {
92 handler
93 .get_table(&table_name, query_ctx)
94 .await
95 .context(CatalogSnafu)?
96 };
97 identity_pipeline(data_array, table, pipeline_ctx)
98 .map(|opt_map| ContextReq::from_opt_map(opt_map, table_name))
99 .context(PipelineSnafu)
100}
101
102async fn run_custom_pipeline(
103 handler: &PipelineHandlerRef,
104 pipeline_ctx: &PipelineContext<'_>,
105 pipeline_req: PipelineIngestRequest,
106 query_ctx: &QueryContextRef,
107 is_top_level: bool,
108) -> Result<ContextReq> {
109 let skip_error = pipeline_ctx.pipeline_param.skip_error();
110 let db = query_ctx.get_db_string();
111 let pipeline = get_pipeline(pipeline_ctx.pipeline_definition, handler, query_ctx).await?;
112
113 let transform_timer = std::time::Instant::now();
114
115 let PipelineIngestRequest {
116 table: table_name,
117 values: pipeline_maps,
118 } = pipeline_req;
119 let arr_len = pipeline_maps.len();
120 let mut transformed_map = HashMap::new();
121 let mut dispatched: BTreeMap<DispatchedTo, Vec<VrlValue>> = BTreeMap::new();
122
123 let mut schema_info = match pipeline.transformer() {
124 TransformerMode::GreptimeTransformer(greptime_transformer) => {
125 SchemaInfo::from_schema_list(greptime_transformer.schemas().clone())
126 }
127 TransformerMode::AutoTransform(ts_name, timeunit) => {
128 let timeunit = match timeunit {
129 TimeUnit::Second => ColumnDataType::TimestampSecond,
130 TimeUnit::Millisecond => ColumnDataType::TimestampMillisecond,
131 TimeUnit::Microsecond => ColumnDataType::TimestampMicrosecond,
132 TimeUnit::Nanosecond => ColumnDataType::TimestampNanosecond,
133 };
134
135 let mut schema_info = SchemaInfo::default();
136 schema_info.schema.push(ColumnSchema {
137 column_name: ts_name.clone(),
138 datatype: timeunit.into(),
139 semantic_type: SemanticType::Timestamp as i32,
140 datatype_extension: None,
141 options: None,
142 });
143
144 schema_info
145 }
146 };
147
148 for pipeline_map in pipeline_maps {
149 let result = pipeline
150 .exec_mut(pipeline_map, pipeline_ctx, &mut schema_info)
151 .inspect_err(|_| {
152 METRIC_HTTP_LOGS_TRANSFORM_ELAPSED
153 .with_label_values(&[db.as_str(), METRIC_FAILURE_VALUE])
154 .observe(transform_timer.elapsed().as_secs_f64());
155 })
156 .context(PipelineSnafu);
157
158 let r = unwrap_or_continue_if_err!(result, skip_error);
159 match r {
160 PipelineExecOutput::Transformed(TransformedOutput {
161 opt,
162 row,
163 table_suffix,
164 }) => {
165 let act_table_name = table_suffix_to_table_name(&table_name, table_suffix);
166 push_to_map!(transformed_map, (opt, act_table_name), row, arr_len);
167 }
168 PipelineExecOutput::DispatchedTo(dispatched_to, val) => {
169 push_to_map!(dispatched, dispatched_to, val, arr_len);
170 }
171 PipelineExecOutput::Filtered => {
172 continue;
173 }
174 }
175 }
176
177 let mut results = ContextReq::default();
178
179 let s_len = schema_info.schema.len();
180
181 for ((opt, table_name), mut rows) in transformed_map {
183 for row in rows.iter_mut() {
184 row.values
185 .resize(s_len, greptime_proto::v1::Value::default());
186 }
187 results.add_row(
188 opt,
189 RowInsertRequest {
190 rows: Some(Rows {
191 rows,
192 schema: schema_info.schema.clone(),
193 }),
194 table_name,
195 },
196 );
197 }
198
199 for (dispatched_to, coll) in dispatched {
202 let table_name = dispatched_to.dispatched_to_table_name(&table_name);
205 let next_pipeline_name = dispatched_to
206 .pipeline
207 .as_deref()
208 .unwrap_or(GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME);
209
210 let next_pipeline_def =
212 PipelineDefinition::from_name(next_pipeline_name, None, None).context(PipelineSnafu)?;
213 let next_pipeline_ctx = PipelineContext::new(
214 &next_pipeline_def,
215 pipeline_ctx.pipeline_param,
216 pipeline_ctx.channel,
217 );
218 let requests = Box::pin(run_pipeline(
219 handler,
220 &next_pipeline_ctx,
221 PipelineIngestRequest {
222 table: table_name,
223 values: coll,
224 },
225 query_ctx,
226 false,
227 ))
228 .await?;
229
230 results.merge(requests);
231 }
232
233 if is_top_level {
234 METRIC_HTTP_LOGS_TRANSFORM_ELAPSED
235 .with_label_values(&[db.as_str(), METRIC_SUCCESS_VALUE])
236 .observe(transform_timer.elapsed().as_secs_f64());
237 }
238
239 Ok(results)
240}
241
242#[inline]
243fn table_suffix_to_table_name(table_name: &String, table_suffix: Option<String>) -> String {
244 match table_suffix {
245 Some(suffix) => format!("{}{}", table_name, suffix),
246 None => table_name.clone(),
247 }
248}