// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements. See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership. The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License. You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied. See the License for the

// specific language governing permissions and limitations

// under the License.

use std::ops::Range;

use std::sync::Arc;

use arrow::array::{

Decimal128Builder, Int32Builder, StringBuilder, StringDictionaryBuilder,

TimestampNanosecondBuilder, UInt16Builder,

};

use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef, TimeUnit};

use arrow::record_batch::RecordBatch;

use rand::prelude::StdRng;

use rand::{Rng, SeedableRng};

#[derive(Debug, Clone)]

struct GeneratorOptions {

row_limit: usize,

pods_per_host: Range<usize>,

containers_per_pod: Range<usize>,

entries_per_container: Range<usize>,

include_nulls: bool,

}

impl Default for GeneratorOptions {

fn default() -> Self {

Self {

row_limit: usize::MAX,

pods_per_host: 1..15,

containers_per_pod: 1..3,

entries_per_container: 1024..8192,

include_nulls: false,

}

}

}

/// Creates access log like entries

#[derive(Default)]

struct BatchBuilder {

service: StringDictionaryBuilder<Int32Type>,

host: StringDictionaryBuilder<Int32Type>,

pod: StringDictionaryBuilder<Int32Type>,

container: StringDictionaryBuilder<Int32Type>,

image: StringDictionaryBuilder<Int32Type>,

time: TimestampNanosecondBuilder,

client_addr: StringBuilder,

request_duration: Int32Builder,

request_user_agent: StringBuilder,

request_method: StringBuilder,

request_host: StringBuilder,

request_bytes: Int32Builder,

response_bytes: Int32Builder,

response_status: UInt16Builder,

prices_status: Decimal128Builder,

options: GeneratorOptions,

row_count: usize,

}

impl BatchBuilder {

fn schema() -> SchemaRef {

let utf8_dict =

|| DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));

Arc::new(Schema::new(vec![

Field::new("service", utf8_dict(), true),

Field::new("host", utf8_dict(), false),

Field::new("pod", utf8_dict(), false),

Field::new("container", utf8_dict(), false),

Field::new("image", utf8_dict(), false),

Field::new(

"time",

DataType::Timestamp(TimeUnit::Nanosecond, None),

false,

),

Field::new("client_addr", DataType::Utf8, true),

Field::new("request_duration_ns", DataType::Int32, false),

Field::new("request_user_agent", DataType::Utf8, true),

Field::new("request_method", DataType::Utf8, true),

Field::new("request_host", DataType::Utf8, true),

Field::new("request_bytes", DataType::Int32, true),

Field::new("response_bytes", DataType::Int32, true),

Field::new("response_status", DataType::UInt16, false),

Field::new("decimal_price", DataType::Decimal128(38, 0), false),

]))

}

fn is_finished(&self) -> bool {

self.options.row_limit <= self.row_count

}

fn append(&mut self, rng: &mut StdRng, host: &str, service: &str) {

let num_pods = rng.random_range(self.options.pods_per_host.clone());

let pods = generate_sorted_strings(rng, num_pods, 30..40);

for pod in pods {

let num_containers =

rng.random_range(self.options.containers_per_pod.clone());

for container_idx in 0..num_containers {

let container = format!("{service}_container_{container_idx}");

let image = format!(

"{container}@sha256:30375999bf03beec2187843017b10c9e88d8b1a91615df4eb6350fb39472edd9"

);

let num_entries =

rng.random_range(self.options.entries_per_container.clone());

for i in 0..num_entries {

if self.is_finished() {

return;

}

let time = i as i64 * 1024;

self.append_row(rng, host, &pod, service, &container, &image, time);

}

}

}

}

#[allow(clippy::too_many_arguments)]

fn append_row(

&mut self,

rng: &mut StdRng,

host: &str,

pod: &str,

service: &str,

container: &str,

image: &str,

time: i64,

) {

self.row_count += 1;

let methods = &["GET", "PUT", "POST", "HEAD", "PATCH", "DELETE"];

let status = &[200, 204, 400, 503, 403];

self.service.append(service).unwrap();

self.host.append(host).unwrap();

self.pod.append(pod).unwrap();

self.container.append(container).unwrap();

self.image.append(image).unwrap();

self.time.append_value(time);

if self.options.include_nulls {

// Append a null value if the option is set

// Use both "NULL" as a string and a null value

if rng.random_bool(0.5) {

self.client_addr.append_null();

} else {

self.client_addr.append_value("NULL");

}

} else {

self.client_addr.append_value(format!(

"{}.{}.{}.{}",

rng.random::<u8>(),

rng.random::<u8>(),

rng.random::<u8>(),

rng.random::<u8>()

));

}

self.request_duration.append_value(rng.random());

self.request_user_agent

.append_value(random_string(rng, 20..100));

self.request_method

.append_value(methods[rng.random_range(0..methods.len())]);

self.request_host

.append_value(format!("https://{service}.mydomain.com"));

self.request_bytes

.append_option(rng.random_bool(0.9).then(|| rng.random()));

self.response_bytes

.append_option(rng.random_bool(0.9).then(|| rng.random()));

self.response_status

.append_value(status[rng.random_range(0..status.len())]);

self.prices_status.append_value(self.row_count as i128);

}

fn finish(mut self, schema: SchemaRef) -> RecordBatch {

RecordBatch::try_new(

schema,

vec![

Arc::new(self.service.finish()),

Arc::new(self.host.finish()),

Arc::new(self.pod.finish()),

Arc::new(self.container.finish()),

Arc::new(self.image.finish()),

Arc::new(self.time.finish()),

Arc::new(self.client_addr.finish()),

Arc::new(self.request_duration.finish()),

Arc::new(self.request_user_agent.finish()),

Arc::new(self.request_method.finish()),

Arc::new(self.request_host.finish()),

Arc::new(self.request_bytes.finish()),

Arc::new(self.response_bytes.finish()),

Arc::new(self.response_status.finish()),

Arc::new(

self.prices_status

.finish()

.with_precision_and_scale(38, 0)

.unwrap(),

),

],

)

.unwrap()

}

}

fn random_string(rng: &mut StdRng, len_range: Range<usize>) -> String {

let len = rng.random_range(len_range);

(0..len)

.map(|_| rng.random_range(b'a'..=b'z') as char)

.collect::<String>()

}

fn generate_sorted_strings(

rng: &mut StdRng,

count: usize,

str_len: Range<usize>,

) -> Vec<String> {

let mut strings: Vec<_> = (0..count)

.map(|_| random_string(rng, str_len.clone()))

.collect();

strings.sort_unstable();

strings

}

/// Iterator that generates sorted, [`RecordBatch`]es with randomly generated data with

/// an access log style schema for tracing or monitoring type

/// use cases.

///

/// This is useful for writing tests queries on such data

///

/// Here are the columns with example data:

///

/// ```text

/// service: 'backend'

/// host: 'i-1ec3ca3151468928.ec2.internal'

/// pod: 'aqcathnxqsphdhgjtgvxsfyiwbmhlmg'

/// container: 'backend_container_0'

/// image: 'backend_container_0@sha256:30375999bf03beec2187843017b10c9e88d8b1a91615df4eb6350fb39472edd9'

/// time: '1970-01-01 00:00:00'

/// client_addr: '127.216.178.64'

/// request_duration_ns: -1261239112

/// request_user_agent: 'kxttrfiiietlsaygzphhwlqcgngnumuphliejmxfdznuurswhdcicrlprbnocibvsbukiohjjbjdygwbfhxqvurm'

/// request_method: 'PUT'

/// request_host: 'https://backend.mydomain.com'

/// request_bytes: -312099516

/// response_bytes: 1448834362

/// response_status: 200

/// ```

#[derive(Debug)]

pub struct AccessLogGenerator {

schema: SchemaRef,

rng: StdRng,

host_idx: usize,

/// maximum rows per batch

max_batch_size: usize,

/// How many rows have been returned so far

row_count: usize,

/// Options

options: GeneratorOptions,

}

impl Default for AccessLogGenerator {

fn default() -> Self {

Self::new()

}

}

impl AccessLogGenerator {

pub fn new() -> Self {

let seed = [

1, 0, 0, 0, 23, 0, 3, 0, 200, 1, 0, 0, 210, 30, 8, 0, 1, 0, 21, 0, 6, 0, 0,

0, 0, 0, 5, 0, 0, 0, 0, 0,

];

Self {

schema: BatchBuilder::schema(),

host_idx: 0,

rng: StdRng::from_seed(seed),

max_batch_size: usize::MAX,

row_count: 0,

options: Default::default(),

}

}

/// Return the schema of the [`RecordBatch`]es created

pub fn schema(&self) -> SchemaRef {

self.schema.clone()

}

/// Limit the maximum batch size

pub fn with_max_batch_size(mut self, batch_size: usize) -> Self {

self.max_batch_size = batch_size;

self

}

/// Return up to row_limit rows;

pub fn with_row_limit(mut self, row_limit: usize) -> Self {

self.options.row_limit = row_limit;

self

}

/// Set the number of pods per host

pub fn with_pods_per_host(mut self, range: Range<usize>) -> Self {

self.options.pods_per_host = range;

self

}

/// Set the number of containers per pod

pub fn with_containers_per_pod(mut self, range: Range<usize>) -> Self {

self.options.containers_per_pod = range;

self

}

/// Set the number of log entries per container

pub fn with_entries_per_container(mut self, range: Range<usize>) -> Self {

self.options.entries_per_container = range;

self

}

// Set the condition for null values in the generated data

pub fn with_include_nulls(mut self, include_nulls: bool) -> Self {

self.options.include_nulls = include_nulls;

self

}

}

impl Iterator for AccessLogGenerator {

type Item = RecordBatch;

fn next(&mut self) -> Option<Self::Item> {

if self.row_count == self.options.row_limit {

return None;

}

let row_limit = self

.max_batch_size

.min(self.options.row_limit - self.row_count);

let mut builder = BatchBuilder {

options: GeneratorOptions {

row_limit,

..self.options.clone()

},

..Default::default()

};

let host = format!(

"i-{:016x}.ec2.internal",

self.host_idx * 0x7d87f8ed5c5 + 0x1ec3ca3151468928

);

self.host_idx += 1;

for service in &["frontend", "backend", "database", "cache"] {

if self.rng.random_bool(0.5) {

continue;

}

if builder.is_finished() {

break;

}

builder.append(&mut self.rng, &host, service);

}

let batch = builder.finish(Arc::clone(&self.schema));

self.row_count += batch.num_rows();

Some(batch)

}

}