import unittest

from feldera import PipelineBuilder

from tests import TEST_CLIENT

from tests.platform.helper import PipelineTestCase

from feldera.runtime_config import Resources, RuntimeConfig

from feldera.testutils import (

FELDERA_TEST_NUM_WORKERS,

FELDERA_TEST_NUM_HOSTS,

min_datafusion_memory_mb,

)

# Test user-defined preprocessor

class TestUDP(PipelineTestCase):

def test_local(self):

sql = """

CREATE TABLE t (

i INT,

ti TINYINT,

si SMALLINT,

bi BIGINT,

r REAL,

d DOUBLE,

bin VARBINARY,

dt DATE,

t TIME,

ts TIMESTAMP,

a INT ARRAY,

b BOOLEAN,

dc DECIMAL(7,2),

s VARCHAR

) WITH ('connectors' = '[{

"name": "t",

"transport": {

"name": "datagen",

"config": {

"seed": 1,

"plan": [{

"limit": 100000

}]

}

},

"preprocessor": [{

"name": "logger",

"message_oriented": true,

"config": {}

}]

}]');

CREATE MATERIALIZED VIEW v AS

SELECT * FROM t;

"""

udfs = """

use tracing::info;

use std::sync::{Arc, Mutex};

use feldera_adapterlib::format::{ParseError, Splitter};

use feldera_adapterlib::preprocess::{

Preprocessor, PreprocessorCreateError, PreprocessorFactory,

};

use feldera_types::preprocess::PreprocessorConfig;

pub struct LoggerPreprocessor {

count: Arc<Mutex<u64>>,

}

impl Preprocessor for LoggerPreprocessor {

fn process(&mut self, data: &[u8]) -> (Vec<u8>, Vec<ParseError>) {

let mut count = self.count.lock().unwrap();

*count += data.len() as u64;

// Log a message if the counter has crossed a Megabyte boundary

if *count / (1024 * 1024) > (*count - data.len() as u64) / (1024 * 1024) {

info!("Processed {} bytes of data", *count);

}

(data.to_vec(), vec![])

}

fn fork(&self) -> Box<dyn Preprocessor> {

Box::new(LoggerPreprocessor { count: Arc::clone(&self.count) })

}

fn splitter(&self) -> Option<Box<dyn Splitter>> {

None

}

}

pub struct LoggerPreprocessorFactory;

impl PreprocessorFactory for LoggerPreprocessorFactory {

fn create(

&self,

_config: &PreprocessorConfig,

) -> Result<Box<dyn Preprocessor>, PreprocessorCreateError> {

Ok(Box::new(LoggerPreprocessor { count: Arc::new(Mutex::new(0)) }))

}

}

"""

toml = """

tracing = { version = "0.1.40" }

"""

# Scales with FELDERA_TEST_NUM_WORKERS/_HOSTS so the ad-hoc ORDER BY

# query below doesn't fail with "Resources exhausted" if either

# changes. See min_datafusion_memory_mb.

datafusion_memory_mb = min_datafusion_memory_mb(

FELDERA_TEST_NUM_WORKERS, FELDERA_TEST_NUM_HOSTS

)

pipeline = PipelineBuilder(

TEST_CLIENT,

name=self.register_for_cleanup("test_udps"),

sql=sql,

udf_rust=udfs,

udf_toml=toml,

runtime_config=RuntimeConfig(

datafusion_memory_mb=datafusion_memory_mb,

resources=Resources(memory_mb_min=datafusion_memory_mb + 512),

workers=FELDERA_TEST_NUM_WORKERS,

hosts=FELDERA_TEST_NUM_HOSTS,

),

).create_or_replace()

pipeline.start_paused()

pipeline.resume()

pipeline.wait_for_completion()

for log in pipeline.logs():

# This will loop forever if the message is not found

if "bytes of data" in log:

break

hash = pipeline.query_hash("SELECT * FROM v ORDER BY i, ti, si, bi")

assert (

hash == "29C468885BA59EAD1A0BA70C07A997401030261AB1CDEB9EA7B13B0DE6722646"

), "Hash does not match"

pipeline.stop(force=True)

pipeline.delete(True)

if __name__ == "__main__":

unittest.main()