1use std::sync::Arc;
2
3use anyhow::{Context, Result};
4use chrono::DateTime;
5use datafusion::arrow::array::ArrayBuilder;
6use datafusion::arrow::array::BinaryDictionaryBuilder;
7use datafusion::arrow::array::PrimitiveBuilder;
8use datafusion::arrow::array::StringBuilder;
9use datafusion::arrow::array::StringDictionaryBuilder;
10use datafusion::arrow::datatypes::DataType;
11use datafusion::arrow::datatypes::Field;
12use datafusion::arrow::datatypes::Int32Type;
13use datafusion::arrow::datatypes::Schema;
14use datafusion::arrow::datatypes::TimeUnit;
15use datafusion::arrow::datatypes::TimestampNanosecondType;
16use datafusion::arrow::record_batch::RecordBatch;
17
18use crate::log_entry::LogEntry;
19use crate::metadata::ProcessMetadata;
20use crate::properties::property_set_jsonb_dictionary_builder::PropertySetJsonbDictionaryBuilder;
21use crate::time::TimeRange;
22
23pub fn log_table_schema() -> Schema {
25 Schema::new(vec![
26 Field::new(
27 "process_id",
28 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
29 false,
30 ),
31 Field::new(
32 "stream_id",
33 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
34 false,
35 ),
36 Field::new(
37 "block_id",
38 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
39 false,
40 ),
41 Field::new(
42 "insert_time",
43 DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
44 false,
45 ),
46 Field::new(
47 "exe",
48 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
49 false,
50 ),
51 Field::new(
52 "username",
53 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
54 false,
55 ),
56 Field::new(
57 "computer",
58 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
59 false,
60 ),
61 Field::new(
62 "time",
63 DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
64 false,
65 ),
66 Field::new(
67 "target",
68 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
69 false,
70 ),
71 Field::new("level", DataType::Int32, false),
72 Field::new("msg", DataType::Utf8, false),
73 Field::new(
74 "properties",
75 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)),
76 false,
77 ),
78 Field::new(
79 "process_properties",
80 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)),
81 false,
82 ),
83 Field::new(
89 "audience",
90 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
91 false,
92 ),
93 ])
94}
95
96pub struct LogEntriesRecordBuilder {
98 process_ids: StringDictionaryBuilder<Int32Type>,
99 stream_ids: StringDictionaryBuilder<Int32Type>,
100 block_ids: StringDictionaryBuilder<Int32Type>,
101 insert_times: PrimitiveBuilder<TimestampNanosecondType>,
102 exes: StringDictionaryBuilder<Int32Type>,
103 usernames: StringDictionaryBuilder<Int32Type>,
104 computers: StringDictionaryBuilder<Int32Type>,
105 times: PrimitiveBuilder<TimestampNanosecondType>,
106 targets: StringDictionaryBuilder<Int32Type>,
107 levels: PrimitiveBuilder<Int32Type>,
108 msgs: StringBuilder,
109 properties: PropertySetJsonbDictionaryBuilder,
110 process_properties: BinaryDictionaryBuilder<Int32Type>,
111 audiences: StringDictionaryBuilder<Int32Type>,
112}
113
114impl LogEntriesRecordBuilder {
115 pub fn with_capacity(capacity: usize) -> Self {
116 Self {
117 process_ids: StringDictionaryBuilder::new(),
118 stream_ids: StringDictionaryBuilder::new(),
119 block_ids: StringDictionaryBuilder::new(),
120 insert_times: PrimitiveBuilder::with_capacity(capacity),
121 exes: StringDictionaryBuilder::new(),
122 usernames: StringDictionaryBuilder::new(),
123 computers: StringDictionaryBuilder::new(),
124 times: PrimitiveBuilder::with_capacity(capacity),
125 targets: StringDictionaryBuilder::new(),
126 levels: PrimitiveBuilder::with_capacity(capacity),
127 msgs: StringBuilder::new(),
128 properties: PropertySetJsonbDictionaryBuilder::new(capacity),
129 process_properties: BinaryDictionaryBuilder::new(),
130 audiences: StringDictionaryBuilder::new(),
131 }
132 }
133
134 pub fn get_time_range(&self) -> Option<TimeRange> {
135 if self.is_empty() {
136 return None;
137 }
138 let slice = self.times.values_slice();
140 Some(TimeRange::new(
141 DateTime::from_timestamp_nanos(slice[0]),
142 DateTime::from_timestamp_nanos(slice[slice.len() - 1]),
143 ))
144 }
145
146 pub fn len(&self) -> i64 {
147 self.times.len() as i64
148 }
149
150 pub fn is_empty(&self) -> bool {
151 self.times.len() == 0
152 }
153
154 pub fn append(&mut self, row: &LogEntry) -> Result<()> {
155 self.process_ids
156 .append(format!("{}", row.process.process_id))?;
157 self.stream_ids.append(&*row.stream_id)?;
158 self.block_ids.append(&*row.block_id)?;
159 self.insert_times.append_value(row.insert_time);
160 self.exes.append(&row.process.exe)?;
161 self.usernames.append(&row.process.username)?;
162 self.computers.append(&row.process.computer)?;
163 self.times.append_value(row.time);
164 self.targets.append(row.target)?;
165 self.levels.append_value(row.level);
166 self.msgs.append_value(row.msg);
167 self.properties.append_property_set(&row.properties)?;
168 self.process_properties.append(&*row.process.properties)?;
169 self.audiences.append(&*row.process.audience)?;
170 Ok(())
171 }
172
173 pub fn append_entry_only(&mut self, row: &LogEntry) -> Result<()> {
175 self.times.append_value(row.time);
177 self.targets.append(row.target)?;
178 self.levels.append_value(row.level);
179 self.msgs.append_value(row.msg);
180 self.properties.append_property_set(&row.properties)?;
181 Ok(())
182 }
183
184 pub fn fill_constant_columns(
186 &mut self,
187 process: &ProcessMetadata,
188 stream_id: &str,
189 block_id: &str,
190 insert_time: i64,
191 entry_count: usize,
192 ) -> Result<()> {
193 let process_id_str = format!("{}", process.process_id);
194
195 let insert_times_slice = vec![insert_time; entry_count];
197 self.insert_times.append_slice(&insert_times_slice);
198
199 self.process_ids.append_n(&process_id_str, entry_count)?;
200 self.stream_ids.append_n(stream_id, entry_count)?;
201 self.block_ids.append_n(block_id, entry_count)?;
202 self.exes.append_n(&process.exe, entry_count)?;
203 self.usernames.append_n(&process.username, entry_count)?;
204 self.computers.append_n(&process.computer, entry_count)?;
205 self.process_properties
206 .append_n(&**process.properties, entry_count)?;
207 self.audiences.append_n(&*process.audience, entry_count)?;
208
209 Ok(())
210 }
211
212 pub fn finish(mut self) -> Result<RecordBatch> {
213 RecordBatch::try_new(
214 Arc::new(log_table_schema()),
215 vec![
216 Arc::new(self.process_ids.finish()),
217 Arc::new(self.stream_ids.finish()),
218 Arc::new(self.block_ids.finish()),
219 Arc::new(self.insert_times.finish().with_timezone_utc()),
220 Arc::new(self.exes.finish()),
221 Arc::new(self.usernames.finish()),
222 Arc::new(self.computers.finish()),
223 Arc::new(self.times.finish().with_timezone_utc()),
224 Arc::new(self.targets.finish()),
225 Arc::new(self.levels.finish()),
226 Arc::new(self.msgs.finish()),
227 Arc::new(self.properties.finish()?),
228 Arc::new(self.process_properties.finish()),
229 Arc::new(self.audiences.finish()),
230 ],
231 )
232 .with_context(|| "building record batch")
233 }
234}