tests_integration/
test_util.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::env;
16use std::fmt::Display;
17use std::net::SocketAddr;
18use std::sync::Arc;
19use std::time::Duration;
20
21use auth::{DefaultPermissionChecker, PermissionCheckerRef, UserProviderRef};
22use axum::Router;
23use catalog::kvbackend::KvBackendCatalogManager;
24use common_base::Plugins;
25use common_config::Configurable;
26use common_meta::key::TableMetadataManager;
27use common_meta::key::catalog_name::CatalogNameKey;
28use common_meta::key::schema_name::SchemaNameKey;
29use common_meta::kv_backend::KvBackendRef;
30use common_query::Output;
31use common_runtime::runtime::BuilderBuild;
32use common_runtime::{Builder as RuntimeBuilder, Runtime};
33use common_test_util::ports;
34use common_test_util::temp_dir::{TempDir, create_temp_dir};
35use common_wal::config::DatanodeWalConfig;
36use datanode::config::{DatanodeOptions, StorageConfig};
37use frontend::instance::Instance;
38use frontend::service_config::{MysqlOptions, PostgresOptions};
39use mito2::gc::GcConfig;
40use object_store::config::{
41    AzblobConfig, FileConfig, GcsConfig, ObjectStoreConfig, OssConfig, S3Config,
42};
43use object_store::services::{Azblob, Gcs, Oss, S3};
44use object_store::test_util::TempFolder;
45use object_store::{AzblobConnection, GcsConnection, ObjectStore, OssConnection, S3Connection};
46use servers::grpc::builder::GrpcServerBuilder;
47use servers::grpc::greptime_handler::GreptimeRequestHandler;
48use servers::grpc::{FlightCompression, GrpcOptions, GrpcServer, GrpcServerConfig};
49use servers::http::{HttpOptions, HttpServerBuilder};
50use servers::metrics_handler::MetricsHandler;
51use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef};
52use servers::otel_arrow::OtelArrowServiceHandler;
53use servers::pending_rows_batcher::PendingRowsBatcher;
54use servers::postgres::PostgresServer;
55use servers::prom_remote_write::validation::PromValidationMode;
56use servers::query_handler::sql::SqlQueryHandler;
57use servers::request_memory_limiter::ServerMemoryLimiter;
58use servers::server::Server;
59use servers::tls::ReloadableTlsServerConfig;
60use session::context::QueryContext;
61
62use crate::standalone::{GreptimeDbStandalone, GreptimeDbStandaloneBuilder};
63
64pub const PEER_PLACEHOLDER_ADDR: &str = "127.0.0.1:3001";
65
66#[derive(Debug, Clone, Copy, Eq, PartialEq)]
67pub enum StorageType {
68    S3,
69    S3WithCache,
70    File,
71    Oss,
72    Azblob,
73    Gcs,
74}
75
76impl Display for StorageType {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        match self {
79            StorageType::S3 => write!(f, "S3"),
80            StorageType::S3WithCache => write!(f, "S3"),
81            StorageType::File => write!(f, "File"),
82            StorageType::Oss => write!(f, "Oss"),
83            StorageType::Azblob => write!(f, "Azblob"),
84            StorageType::Gcs => write!(f, "Gcs"),
85        }
86    }
87}
88
89impl StorageType {
90    pub fn build_storage_types_based_on_env() -> Vec<StorageType> {
91        let mut storage_types = Vec::with_capacity(4);
92        storage_types.push(StorageType::File);
93        if let Ok(bucket) = env::var("GT_S3_BUCKET")
94            && !bucket.is_empty()
95        {
96            storage_types.push(StorageType::S3);
97        }
98        if env::var("GT_OSS_BUCKET").is_ok() {
99            storage_types.push(StorageType::Oss);
100        }
101        if env::var("GT_AZBLOB_CONTAINER").is_ok() {
102            storage_types.push(StorageType::Azblob);
103        }
104        if env::var("GT_GCS_BUCKET").is_ok() {
105            storage_types.push(StorageType::Gcs);
106        }
107        storage_types
108    }
109
110    pub fn test_on(&self) -> bool {
111        let _ = dotenv::dotenv();
112
113        match self {
114            StorageType::File => true, // always test file
115            StorageType::S3 | StorageType::S3WithCache => {
116                if let Ok(b) = env::var("GT_S3_BUCKET") {
117                    !b.is_empty()
118                } else {
119                    false
120                }
121            }
122            StorageType::Oss => {
123                if let Ok(b) = env::var("GT_OSS_BUCKET") {
124                    !b.is_empty()
125                } else {
126                    false
127                }
128            }
129            StorageType::Azblob => {
130                if let Ok(b) = env::var("GT_AZBLOB_CONTAINER") {
131                    !b.is_empty()
132                } else {
133                    false
134                }
135            }
136            StorageType::Gcs => {
137                if let Ok(b) = env::var("GT_GCS_BUCKET") {
138                    !b.is_empty()
139                } else {
140                    false
141                }
142            }
143        }
144    }
145}
146
147fn s3_test_config() -> S3Config {
148    S3Config {
149        connection: S3Connection {
150            root: uuid::Uuid::new_v4().to_string(),
151            access_key_id: env::var("GT_S3_ACCESS_KEY_ID").unwrap().into(),
152            secret_access_key: env::var("GT_S3_ACCESS_KEY").unwrap().into(),
153            bucket: env::var("GT_S3_BUCKET").unwrap(),
154            region: Some(env::var("GT_S3_REGION").unwrap()),
155            endpoint: env::var("GT_S3_ENDPOINT_URL").ok(),
156            ..Default::default()
157        },
158        ..Default::default()
159    }
160}
161
162pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, TempDirGuard) {
163    let _ = dotenv::dotenv();
164
165    match store_type {
166        StorageType::Gcs => {
167            let gcs_config = GcsConfig {
168                connection: GcsConnection {
169                    root: uuid::Uuid::new_v4().to_string(),
170                    bucket: env::var("GT_GCS_BUCKET").unwrap(),
171                    scope: env::var("GT_GCS_SCOPE").unwrap(),
172                    credential_path: env::var("GT_GCS_CREDENTIAL_PATH").unwrap().into(),
173                    credential: env::var("GT_GCS_CREDENTIAL").unwrap().into(),
174                    endpoint: env::var("GT_GCS_ENDPOINT").unwrap_or_default(),
175                },
176                ..Default::default()
177            };
178
179            let builder = Gcs::from(&gcs_config.connection);
180            let config = ObjectStoreConfig::Gcs(gcs_config);
181            let store = ObjectStore::new(builder).unwrap().finish();
182            (config, TempDirGuard::Gcs(TempFolder::new(&store, "/")))
183        }
184        StorageType::Azblob => {
185            let azblob_config = AzblobConfig {
186                connection: AzblobConnection {
187                    root: uuid::Uuid::new_v4().to_string(),
188                    container: env::var("GT_AZBLOB_CONTAINER").unwrap(),
189                    account_name: env::var("GT_AZBLOB_ACCOUNT_NAME").unwrap().into(),
190                    account_key: env::var("GT_AZBLOB_ACCOUNT_KEY").unwrap().into(),
191                    endpoint: env::var("GT_AZBLOB_ENDPOINT").unwrap(),
192                    ..Default::default()
193                },
194                ..Default::default()
195            };
196
197            let builder = Azblob::from(&azblob_config.connection);
198            let config = ObjectStoreConfig::Azblob(azblob_config);
199            let store = ObjectStore::new(builder).unwrap().finish();
200            (config, TempDirGuard::Azblob(TempFolder::new(&store, "/")))
201        }
202        StorageType::Oss => {
203            let oss_config = OssConfig {
204                connection: OssConnection {
205                    root: uuid::Uuid::new_v4().to_string(),
206                    access_key_id: env::var("GT_OSS_ACCESS_KEY_ID").unwrap().into(),
207                    access_key_secret: env::var("GT_OSS_ACCESS_KEY").unwrap().into(),
208                    bucket: env::var("GT_OSS_BUCKET").unwrap(),
209                    endpoint: env::var("GT_OSS_ENDPOINT").unwrap(),
210                },
211                ..Default::default()
212            };
213
214            let builder = Oss::from(&oss_config.connection);
215            let config = ObjectStoreConfig::Oss(oss_config);
216            let store = ObjectStore::new(builder).unwrap().finish();
217            (config, TempDirGuard::Oss(TempFolder::new(&store, "/")))
218        }
219        StorageType::S3 | StorageType::S3WithCache => {
220            let mut s3_config = s3_test_config();
221
222            if *store_type == StorageType::S3WithCache {
223                s3_config.cache.cache_path = "/tmp/greptimedb_cache".to_string();
224            } else {
225                s3_config.cache.enable_read_cache = false;
226            }
227
228            let builder = S3::from(&s3_config.connection);
229            let config = ObjectStoreConfig::S3(s3_config);
230            let store = ObjectStore::new(builder).unwrap().finish();
231            (config, TempDirGuard::S3(TempFolder::new(&store, "/")))
232        }
233        StorageType::File => (ObjectStoreConfig::File(FileConfig {}), TempDirGuard::None),
234    }
235}
236
237pub enum TempDirGuard {
238    None,
239    S3(TempFolder),
240    Oss(TempFolder),
241    Azblob(TempFolder),
242    Gcs(TempFolder),
243}
244
245pub struct TestGuard {
246    pub home_guard: FileDirGuard,
247    pub storage_guards: Vec<StorageGuard>,
248}
249
250pub struct FileDirGuard {
251    pub temp_dir: TempDir,
252}
253
254impl FileDirGuard {
255    pub fn new(temp_dir: TempDir) -> Self {
256        Self { temp_dir }
257    }
258}
259
260pub struct StorageGuard(pub TempDirGuard);
261
262impl TestGuard {
263    pub async fn remove_all(&mut self) {
264        for storage_guard in self.storage_guards.iter_mut() {
265            if let TempDirGuard::S3(guard)
266            | TempDirGuard::Oss(guard)
267            | TempDirGuard::Azblob(guard)
268            | TempDirGuard::Gcs(guard) = &mut storage_guard.0
269            {
270                guard.remove_all().await.unwrap()
271            }
272        }
273    }
274}
275
276impl Drop for TestGuard {
277    fn drop(&mut self) {
278        let (tx, rx) = std::sync::mpsc::channel();
279
280        let guards = std::mem::take(&mut self.storage_guards);
281        common_runtime::spawn_global(async move {
282            let mut errors = vec![];
283            for guard in guards {
284                if let TempDirGuard::S3(guard)
285                | TempDirGuard::Oss(guard)
286                | TempDirGuard::Azblob(guard)
287                | TempDirGuard::Gcs(guard) = guard.0
288                    && let Err(e) = guard.remove_all().await
289                {
290                    errors.push(e);
291                }
292            }
293            if errors.is_empty() {
294                tx.send(Ok(())).unwrap();
295            } else {
296                tx.send(Err(errors)).unwrap();
297            }
298        });
299        rx.recv().unwrap().unwrap_or_else(|e| panic!("{:?}", e));
300    }
301}
302
303pub fn create_tmp_dir_and_datanode_opts(
304    default_store_type: StorageType,
305    store_provider_types: Vec<StorageType>,
306    name: &str,
307    wal_config: DatanodeWalConfig,
308    gc_config: GcConfig,
309) -> (DatanodeOptions, TestGuard) {
310    let home_tmp_dir = create_temp_dir(&format!("gt_data_{name}"));
311    let home_dir = home_tmp_dir.path().to_str().unwrap().to_string();
312
313    // Excludes the default object store.
314    let mut store_providers = Vec::with_capacity(store_provider_types.len());
315    // Includes the default object store.
316    let mut storage_guards = Vec::with_capacity(store_provider_types.len() + 1);
317
318    let (default_store, data_tmp_dir) = get_test_store_config(&default_store_type);
319    storage_guards.push(StorageGuard(data_tmp_dir));
320
321    for store_type in store_provider_types {
322        let (store, data_tmp_dir) = get_test_store_config(&store_type);
323        store_providers.push(store);
324        storage_guards.push(StorageGuard(data_tmp_dir))
325    }
326    let opts = create_datanode_opts(
327        default_store,
328        store_providers,
329        home_dir,
330        wal_config,
331        gc_config,
332    );
333
334    (
335        opts,
336        TestGuard {
337            home_guard: FileDirGuard::new(home_tmp_dir),
338            storage_guards,
339        },
340    )
341}
342
343pub(crate) fn create_datanode_opts(
344    default_store: ObjectStoreConfig,
345    providers: Vec<ObjectStoreConfig>,
346    home_dir: String,
347    wal_config: DatanodeWalConfig,
348    gc_config: GcConfig,
349) -> DatanodeOptions {
350    let region_engine = DatanodeOptions::default()
351        .region_engine
352        .into_iter()
353        .map(|mut v| {
354            if let datanode::config::RegionEngineConfig::Mito(mito_config) = &mut v {
355                mito_config.gc = gc_config.clone();
356            }
357            v
358        })
359        .collect();
360    DatanodeOptions {
361        node_id: Some(0),
362        require_lease_before_startup: true,
363        storage: StorageConfig {
364            data_home: home_dir,
365            copy_root: None,
366            providers,
367            store: default_store,
368        },
369        grpc: GrpcOptions::default()
370            .with_bind_addr(PEER_PLACEHOLDER_ADDR)
371            .with_server_addr(PEER_PLACEHOLDER_ADDR),
372        wal: wal_config,
373        region_engine,
374        ..Default::default()
375    }
376}
377
378pub(crate) async fn create_test_table(instance: &Instance, table_name: &str) {
379    let sql = format!(
380        r#"
381CREATE TABLE IF NOT EXISTS {table_name} (
382    host String NOT NULL PRIMARY KEY,
383    cpu DOUBLE NULL,
384    memory DOUBLE NULL,
385    ts TIMESTAMP NOT NULL TIME INDEX,
386)
387"#
388    );
389
390    let result = instance.do_query(&sql, QueryContext::arc()).await;
391    let _ = result.first().unwrap().as_ref().unwrap();
392}
393
394async fn setup_standalone_instance(
395    test_name: &str,
396    store_type: StorageType,
397) -> GreptimeDbStandalone {
398    GreptimeDbStandaloneBuilder::new(test_name)
399        .with_default_store_type(store_type)
400        .build()
401        .await
402}
403
404async fn setup_standalone_instance_with_slow_query_threshold(
405    test_name: &str,
406    store_type: StorageType,
407    slow_query_threshold: std::time::Duration,
408) -> GreptimeDbStandalone {
409    GreptimeDbStandaloneBuilder::new(test_name)
410        .with_default_store_type(store_type)
411        .with_slow_query_threshold(slow_query_threshold)
412        .build()
413        .await
414}
415
416async fn setup_standalone_instance_with_plugins(
417    test_name: &str,
418    store_type: StorageType,
419    plugins: Plugins,
420) -> GreptimeDbStandalone {
421    GreptimeDbStandaloneBuilder::new(test_name)
422        .with_default_store_type(store_type)
423        .with_plugin(plugins)
424        .build()
425        .await
426}
427
428async fn setup_standalone_instance_with_plugins_and_slow_query_threshold(
429    test_name: &str,
430    store_type: StorageType,
431    plugins: Plugins,
432    slow_query_threshold: std::time::Duration,
433) -> GreptimeDbStandalone {
434    GreptimeDbStandaloneBuilder::new(test_name)
435        .with_default_store_type(store_type)
436        .with_plugin(plugins)
437        .with_slow_query_threshold(slow_query_threshold)
438        .build()
439        .await
440}
441
442pub async fn setup_test_http_app(store_type: StorageType, name: &str) -> (Router, TestGuard) {
443    let instance = setup_standalone_instance(name, store_type).await;
444
445    let http_opts = HttpOptions {
446        addr: format!("127.0.0.1:{}", ports::get_port()),
447        ..Default::default()
448    };
449    let http_server = HttpServerBuilder::new(http_opts)
450        .with_sql_handler(instance.fe_instance().clone())
451        .with_logs_handler(instance.fe_instance().clone())
452        .with_metrics_handler(MetricsHandler)
453        .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap())
454        .build();
455    (
456        http_server.build(http_server.make_app()).unwrap(),
457        instance.guard,
458    )
459}
460
461pub async fn setup_test_http_app_with_frontend(
462    store_type: StorageType,
463    name: &str,
464) -> (Router, TestGuard) {
465    setup_test_http_app_with_frontend_and_user_provider(store_type, name, None).await
466}
467
468pub async fn setup_test_http_app_with_frontend_and_slow_query_threshold(
469    store_type: StorageType,
470    name: &str,
471    slow_query_threshold: std::time::Duration,
472) -> (Router, TestGuard) {
473    let instance =
474        setup_standalone_instance_with_slow_query_threshold(name, store_type, slow_query_threshold)
475            .await;
476
477    create_test_table(instance.fe_instance(), "demo").await;
478
479    let http_opts = HttpOptions {
480        addr: format!("127.0.0.1:{}", ports::get_port()),
481        ..Default::default()
482    };
483
484    let http_server = HttpServerBuilder::new(http_opts)
485        .with_sql_handler(instance.fe_instance().clone())
486        .with_log_ingest_handler(instance.fe_instance().clone(), None, None)
487        .with_logs_handler(instance.fe_instance().clone())
488        .with_influxdb_handler(instance.fe_instance().clone())
489        .with_otlp_handler(instance.fe_instance().clone(), true)
490        .with_jaeger_handler(instance.fe_instance().clone())
491        .with_greptime_config_options(instance.opts.to_toml().unwrap())
492        .build();
493
494    let app = http_server.build(http_server.make_app()).unwrap();
495    (app, instance.guard)
496}
497
498pub async fn setup_test_http_app_with_frontend_and_user_provider(
499    store_type: StorageType,
500    name: &str,
501    user_provider: Option<UserProviderRef>,
502) -> (Router, TestGuard) {
503    setup_test_http_app_with_frontend_and_custom_options(
504        store_type,
505        name,
506        user_provider,
507        None,
508        None,
509    )
510    .await
511}
512
513pub async fn setup_test_http_app_with_frontend_and_custom_options(
514    store_type: StorageType,
515    name: &str,
516    user_provider: Option<UserProviderRef>,
517    http_opts: Option<HttpOptions>,
518    memory_limiter: Option<ServerMemoryLimiter>,
519) -> (Router, TestGuard) {
520    let plugins = Plugins::new();
521    if let Some(user_provider) = user_provider.clone() {
522        plugins.insert::<UserProviderRef>(user_provider.clone());
523        plugins.insert::<PermissionCheckerRef>(DefaultPermissionChecker::arc());
524    }
525
526    let instance = setup_standalone_instance_with_plugins(name, store_type, plugins).await;
527
528    create_test_table(instance.fe_instance(), "demo").await;
529
530    let http_opts = http_opts.unwrap_or_else(|| HttpOptions {
531        addr: format!("127.0.0.1:{}", ports::get_port()),
532        ..Default::default()
533    });
534
535    let mut http_server = HttpServerBuilder::new(http_opts)
536        .with_sql_handler(instance.fe_instance().clone())
537        .with_log_ingest_handler(instance.fe_instance().clone(), None, None)
538        .with_logs_handler(instance.fe_instance().clone())
539        .with_influxdb_handler(instance.fe_instance().clone())
540        .with_otlp_handler(instance.fe_instance().clone(), true)
541        .with_jaeger_handler(instance.fe_instance().clone())
542        .with_dashboard_handler(instance.fe_instance().clone())
543        .with_greptime_config_options(instance.opts.to_toml().unwrap());
544
545    if let Some(user_provider) = user_provider {
546        http_server = http_server.with_user_provider(user_provider);
547    }
548
549    if let Some(limiter) = memory_limiter {
550        http_server = http_server.with_memory_limiter(limiter);
551    }
552
553    let http_server = http_server.build();
554
555    let app = http_server.build(http_server.make_app()).unwrap();
556    (app, instance.guard)
557}
558
559async fn run_sql(sql: &str, instance: &GreptimeDbStandalone) {
560    let result = instance
561        .fe_instance()
562        .do_query(sql, QueryContext::arc())
563        .await;
564    let _ = result.first().unwrap().as_ref().unwrap();
565}
566
567pub async fn setup_test_prom_app_with_frontend(
568    store_type: StorageType,
569    name: &str,
570) -> (Router, TestGuard) {
571    setup_test_prom_app_with_frontend_inner(store_type, name, false, false).await
572}
573
574pub async fn setup_test_prom_app_with_frontend_native_histogram(
575    store_type: StorageType,
576    name: &str,
577) -> (Router, TestGuard) {
578    setup_test_prom_app_with_frontend_inner(store_type, name, false, true).await
579}
580
581/// Like [`setup_test_prom_app_with_frontend`] but enables the pending-rows batcher,
582/// so Prometheus remote write goes through the batched (metric-engine) path instead
583/// of the direct `PromStoreProtocolHandler::write` path.
584pub async fn setup_test_prom_app_with_frontend_batched(
585    store_type: StorageType,
586    name: &str,
587) -> (Router, TestGuard) {
588    setup_test_prom_app_with_frontend_inner(store_type, name, true, false).await
589}
590
591async fn setup_test_prom_app_with_frontend_inner(
592    store_type: StorageType,
593    name: &str,
594    enable_batcher: bool,
595    experimental_enable_prometheus_native_histogram: bool,
596) -> (Router, TestGuard) {
597    unsafe {
598        std::env::set_var("TZ", "UTC");
599    }
600
601    let instance = setup_standalone_instance(name, store_type).await;
602
603    // build physical table
604    let sql = "CREATE TABLE phy (ts timestamp time index, val double, host string primary key) engine=metric with ('physical_metric_table' = '')";
605    run_sql(sql, &instance).await;
606    let sql = "CREATE TABLE phy_ns (ts timestamp(0) time index, val double, host string primary key) engine=metric with ('physical_metric_table' = '')";
607    run_sql(sql, &instance).await;
608    // build metric tables
609    let sql = "CREATE TABLE demo (ts timestamp time index, val double, host string primary key) engine=metric with ('on_physical_table' = 'phy')";
610    run_sql(sql, &instance).await;
611    let sql = "CREATE TABLE demo_metrics (ts timestamp time index, val double, idc string primary key) engine=metric with ('on_physical_table' = 'phy')";
612    run_sql(sql, &instance).await;
613    let sql = "CREATE TABLE multi_labels (ts timestamp(0) time index, val double, idc string, env string, host string, primary key (idc, env, host)) engine=metric with ('on_physical_table' = 'phy_ns')";
614    run_sql(sql, &instance).await;
615
616    // insert rows
617    let sql = "INSERT INTO demo(host, val, ts) VALUES ('host1', 1.1, 0), ('host2', 2.1, 600000)";
618    run_sql(sql, &instance).await;
619    let sql =
620        "INSERT INTO demo_metrics(idc, val, ts) VALUES ('idc1', 1.1, 0), ('idc2', 2.1, 600000)";
621    run_sql(sql, &instance).await;
622    // insert a row with empty label
623    let sql = "INSERT INTO demo_metrics(val, ts) VALUES (1.1, 0)";
624    run_sql(sql, &instance).await;
625
626    // insert rows to multi_labels
627    let sql = "INSERT INTO multi_labels(idc, env, host, val, ts) VALUES ('idc1', 'dev', 'host1', 1.1, 0), ('idc1', 'dev', 'host2', 2.1, 0), ('idc2', 'dev', 'host1', 1.1, 0), ('idc2', 'test', 'host3', 2.1, 0)";
628    run_sql(sql, &instance).await;
629
630    // build physical table
631    let sql = "CREATE TABLE phy2 (ts timestamp(9) time index, val double, host string primary key) engine=metric with ('physical_metric_table' = '')";
632    run_sql(sql, &instance).await;
633    let sql = "CREATE TABLE demo_metrics_with_nanos(ts timestamp(9) time index, val double, idc string primary key) engine=metric with ('on_physical_table' = 'phy2')";
634    run_sql(sql, &instance).await;
635    let sql = "INSERT INTO demo_metrics_with_nanos(idc, val, ts) VALUES ('idc1', 1.1, 0)";
636    run_sql(sql, &instance).await;
637
638    // a mito table with non-prometheus compatible values
639    let sql = "CREATE TABLE mito (ts timestamp(9) time index, val double, host bigint primary key) engine=mito";
640    run_sql(sql, &instance).await;
641    let sql = "INSERT INTO mito(host, val, ts) VALUES (1, 1.1, 0)";
642    run_sql(sql, &instance).await;
643
644    let http_opts = HttpOptions {
645        addr: format!("127.0.0.1:{}", ports::get_port()),
646        ..Default::default()
647    };
648    let frontend_ref = instance.fe_instance().clone();
649    // Mirror the production wiring at `frontend::server`: build the batcher from the
650    // instance's managers. A short flush interval keeps the test responsive.
651    let pending_rows_batcher = if enable_batcher {
652        PendingRowsBatcher::try_new(
653            frontend_ref.partition_manager().clone(),
654            frontend_ref.node_manager().clone(),
655            frontend_ref.catalog_manager().clone(),
656            frontend_ref.table_flownode_set_cache().clone(),
657            true,
658            frontend_ref.clone(),
659            Duration::from_millis(50),
660            1000,
661            4,
662            64,
663            64,
664            std::num::NonZeroUsize::new(1024).unwrap(),
665        )
666    } else {
667        None
668    };
669    let http_server = HttpServerBuilder::new(http_opts)
670        .with_sql_handler(frontend_ref.clone())
671        .with_logs_handler(instance.fe_instance().clone())
672        .with_prom_handler(
673            frontend_ref.clone(),
674            Some(frontend_ref.clone()),
675            true,
676            PromValidationMode::Strict,
677            experimental_enable_prometheus_native_histogram,
678            pending_rows_batcher,
679        )
680        .with_prometheus_handler(frontend_ref)
681        .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap())
682        .build();
683    let app = http_server.build(http_server.make_app()).unwrap();
684    (app, instance.guard)
685}
686
687pub async fn setup_grpc_server(
688    store_type: StorageType,
689    name: &str,
690) -> (GreptimeDbStandalone, Arc<GrpcServer>) {
691    setup_grpc_server_with(store_type, name, None, None, None).await
692}
693
694pub async fn setup_grpc_server_with_user_provider(
695    store_type: StorageType,
696    name: &str,
697    user_provider: Option<UserProviderRef>,
698) -> (GreptimeDbStandalone, Arc<GrpcServer>) {
699    setup_grpc_server_with(store_type, name, user_provider, None, None).await
700}
701
702/// Sets up a gRPC server backed by a standalone instance whose frontend has auto
703/// table creation disabled, for testing the server-side global switch.
704pub async fn setup_grpc_server_with_auto_create_table_disabled(
705    store_type: StorageType,
706    name: &str,
707) -> (GreptimeDbStandalone, Arc<GrpcServer>) {
708    let instance = GreptimeDbStandaloneBuilder::new(name)
709        .with_default_store_type(store_type)
710        .with_auto_create_table(false)
711        .build()
712        .await;
713    setup_grpc_server_for_instance(instance, None, None, None).await
714}
715
716pub async fn setup_grpc_server_with(
717    store_type: StorageType,
718    name: &str,
719    user_provider: Option<UserProviderRef>,
720    grpc_config: Option<GrpcServerConfig>,
721    memory_limiter: Option<servers::request_memory_limiter::ServerMemoryLimiter>,
722) -> (GreptimeDbStandalone, Arc<GrpcServer>) {
723    let instance = setup_standalone_instance(name, store_type).await;
724    setup_grpc_server_for_instance(instance, user_provider, grpc_config, memory_limiter).await
725}
726
727/// Builds and starts a gRPC server on top of an already-constructed standalone
728/// instance. This is the shared core behind the `setup_grpc_server_*` helpers.
729async fn setup_grpc_server_for_instance(
730    instance: GreptimeDbStandalone,
731    user_provider: Option<UserProviderRef>,
732    grpc_config: Option<GrpcServerConfig>,
733    memory_limiter: Option<servers::request_memory_limiter::ServerMemoryLimiter>,
734) -> (GreptimeDbStandalone, Arc<GrpcServer>) {
735    let runtime: Runtime = RuntimeBuilder::default()
736        .worker_threads(2)
737        .thread_name("grpc-handlers")
738        .build()
739        .unwrap();
740
741    let fe_instance_ref = instance.fe_instance().clone();
742
743    let greptime_request_handler = GreptimeRequestHandler::new(
744        fe_instance_ref.clone(),
745        user_provider.clone(),
746        Some(runtime.clone()),
747        FlightCompression::default(),
748    );
749
750    let flight_handler = Arc::new(greptime_request_handler.clone());
751
752    let grpc_config = grpc_config.unwrap_or_default();
753    let mut grpc_builder = GrpcServerBuilder::new(grpc_config.clone(), runtime);
754
755    if let Some(limiter) = memory_limiter {
756        grpc_builder = grpc_builder.with_memory_limiter(limiter);
757    }
758
759    let grpc_builder = grpc_builder
760        .database_handler(greptime_request_handler)
761        .flight_handler(flight_handler)
762        .prometheus_handler(fe_instance_ref.clone(), user_provider.clone())
763        .otel_arrow_handler(OtelArrowServiceHandler::new(fe_instance_ref, user_provider))
764        .with_tls_config(grpc_config.tls)
765        .unwrap();
766
767    let mut grpc_server = grpc_builder.build();
768
769    let fe_grpc_addr = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
770    grpc_server.start(fe_grpc_addr).await.unwrap();
771
772    (instance, Arc::new(grpc_server))
773}
774
775pub async fn setup_mysql_server(
776    store_type: StorageType,
777    name: &str,
778) -> (TestGuard, Arc<Box<dyn Server>>) {
779    setup_mysql_server_with_user_provider(store_type, name, None).await
780}
781
782pub async fn setup_mysql_server_with_slow_query_threshold(
783    store_type: StorageType,
784    name: &str,
785    slow_query_threshold: std::time::Duration,
786) -> (TestGuard, Arc<Box<dyn Server>>) {
787    let plugins = Plugins::new();
788    let instance = setup_standalone_instance_with_plugins_and_slow_query_threshold(
789        name,
790        store_type,
791        plugins,
792        slow_query_threshold,
793    )
794    .await;
795
796    let runtime = RuntimeBuilder::default()
797        .worker_threads(2)
798        .thread_name("mysql-runtime")
799        .build()
800        .unwrap();
801
802    let fe_mysql_addr = format!("127.0.0.1:{}", ports::get_port());
803
804    let fe_instance_ref = instance.fe_instance().clone();
805    let opts = MysqlOptions {
806        addr: fe_mysql_addr.clone(),
807        ..Default::default()
808    };
809    let mut mysql_server = MysqlServer::create_server(
810        runtime,
811        Arc::new(MysqlSpawnRef::new(fe_instance_ref, None)),
812        Arc::new(MysqlSpawnConfig::new(
813            false,
814            Arc::new(
815                ReloadableTlsServerConfig::try_new(opts.tls.clone())
816                    .expect("Failed to load certificates and keys"),
817            ),
818            0,
819            opts.reject_no_database.unwrap_or(false),
820            opts.prepared_stmt_cache_size,
821        )),
822        None,
823    );
824
825    mysql_server
826        .start(fe_mysql_addr.parse::<SocketAddr>().unwrap())
827        .await
828        .unwrap();
829
830    (instance.guard, Arc::new(mysql_server))
831}
832
833pub async fn setup_mysql_server_with_user_provider(
834    store_type: StorageType,
835    name: &str,
836    user_provider: Option<UserProviderRef>,
837) -> (TestGuard, Arc<Box<dyn Server>>) {
838    let plugins = Plugins::new();
839    if let Some(user_provider) = user_provider.clone() {
840        plugins.insert::<UserProviderRef>(user_provider.clone());
841        plugins.insert::<PermissionCheckerRef>(DefaultPermissionChecker::arc());
842    }
843
844    let instance = setup_standalone_instance_with_plugins(name, store_type, plugins).await;
845
846    let runtime = RuntimeBuilder::default()
847        .worker_threads(2)
848        .thread_name("mysql-runtime")
849        .build()
850        .unwrap();
851
852    let fe_mysql_addr = format!("127.0.0.1:{}", ports::get_port());
853
854    let fe_instance_ref = instance.fe_instance().clone();
855    let opts = MysqlOptions {
856        addr: fe_mysql_addr.clone(),
857        ..Default::default()
858    };
859    let mut mysql_server = MysqlServer::create_server(
860        runtime,
861        Arc::new(MysqlSpawnRef::new(fe_instance_ref, user_provider)),
862        Arc::new(MysqlSpawnConfig::new(
863            false,
864            Arc::new(
865                ReloadableTlsServerConfig::try_new(opts.tls.clone())
866                    .expect("Failed to load certificates and keys"),
867            ),
868            0,
869            opts.reject_no_database.unwrap_or(false),
870            opts.prepared_stmt_cache_size,
871        )),
872        None,
873    );
874
875    mysql_server
876        .start(fe_mysql_addr.parse::<SocketAddr>().unwrap())
877        .await
878        .unwrap();
879
880    (instance.guard, Arc::new(mysql_server))
881}
882
883pub async fn setup_pg_server(
884    store_type: StorageType,
885    name: &str,
886) -> (TestGuard, Arc<Box<dyn Server>>) {
887    setup_pg_server_with_user_provider(store_type, name, None).await
888}
889
890pub async fn setup_pg_server_with_slow_query_threshold(
891    store_type: StorageType,
892    name: &str,
893    slow_query_threshold: std::time::Duration,
894) -> (TestGuard, Arc<Box<dyn Server>>) {
895    let instance =
896        setup_standalone_instance_with_slow_query_threshold(name, store_type, slow_query_threshold)
897            .await;
898
899    let runtime = RuntimeBuilder::default()
900        .worker_threads(2)
901        .thread_name("pg-runtime")
902        .build()
903        .unwrap();
904
905    let fe_pg_addr = format!("127.0.0.1:{}", ports::get_port());
906
907    let fe_instance_ref = instance.fe_instance().clone();
908    let opts = PostgresOptions {
909        addr: fe_pg_addr.clone(),
910        ..Default::default()
911    };
912    let tls_server_config = Arc::new(
913        ReloadableTlsServerConfig::try_new(opts.tls.clone())
914            .expect("Failed to load certificates and keys"),
915    );
916
917    let mut pg_server = Box::new(PostgresServer::new(
918        fe_instance_ref,
919        opts.tls.should_force_tls(),
920        tls_server_config,
921        0,
922        runtime,
923        None,
924        None,
925    ));
926
927    pg_server
928        .start(fe_pg_addr.parse::<SocketAddr>().unwrap())
929        .await
930        .unwrap();
931
932    (instance.guard, Arc::new(pg_server))
933}
934
935pub async fn setup_pg_server_with_user_provider(
936    store_type: StorageType,
937    name: &str,
938    user_provider: Option<UserProviderRef>,
939) -> (TestGuard, Arc<Box<dyn Server>>) {
940    let instance = setup_standalone_instance(name, store_type).await;
941
942    let runtime = RuntimeBuilder::default()
943        .worker_threads(2)
944        .thread_name("pg-runtime")
945        .build()
946        .unwrap();
947
948    let fe_pg_addr = format!("127.0.0.1:{}", ports::get_port());
949
950    let fe_instance_ref = instance.fe_instance().clone();
951    let opts = PostgresOptions {
952        addr: fe_pg_addr.clone(),
953        ..Default::default()
954    };
955    let tls_server_config = Arc::new(
956        ReloadableTlsServerConfig::try_new(opts.tls.clone())
957            .expect("Failed to load certificates and keys"),
958    );
959
960    let mut pg_server = Box::new(PostgresServer::new(
961        fe_instance_ref,
962        opts.tls.should_force_tls(),
963        tls_server_config,
964        0,
965        runtime,
966        user_provider,
967        None,
968    ));
969
970    pg_server
971        .start(fe_pg_addr.parse::<SocketAddr>().unwrap())
972        .await
973        .unwrap();
974
975    (instance.guard, Arc::new(pg_server))
976}
977
978pub(crate) async fn prepare_another_catalog_and_schema(instance: &Instance) {
979    let catalog_manager = instance
980        .catalog_manager()
981        .as_any()
982        .downcast_ref::<KvBackendCatalogManager>()
983        .unwrap();
984
985    let table_metadata_manager = catalog_manager.table_metadata_manager_ref();
986    prepare_another_catalog_and_schema_with_manager(table_metadata_manager).await;
987}
988
989pub(crate) async fn prepare_another_catalog_and_schema_with_kv_backend(kv_backend: KvBackendRef) {
990    let table_metadata_manager = TableMetadataManager::new(kv_backend);
991    prepare_another_catalog_and_schema_with_manager(&table_metadata_manager).await;
992}
993
994async fn prepare_another_catalog_and_schema_with_manager(
995    table_metadata_manager: &TableMetadataManager,
996) {
997    table_metadata_manager
998        .catalog_manager()
999        .create(CatalogNameKey::new("another_catalog"), true)
1000        .await
1001        .unwrap();
1002    table_metadata_manager
1003        .schema_manager()
1004        .create(
1005            SchemaNameKey::new("another_catalog", "another_schema"),
1006            None,
1007            true,
1008        )
1009        .await
1010        .unwrap();
1011}
1012
1013pub async fn execute_sql(instance: &Arc<Instance>, sql: &str) -> Output {
1014    SqlQueryHandler::do_query(instance.as_ref(), sql, QueryContext::arc())
1015        .await
1016        .remove(0)
1017        .unwrap()
1018}
1019
1020pub async fn try_execute_sql(
1021    instance: &Arc<Instance>,
1022    sql: &str,
1023) -> servers::error::Result<Output> {
1024    SqlQueryHandler::do_query(instance.as_ref(), sql, QueryContext::arc())
1025        .await
1026        .remove(0)
1027}
1028
1029pub async fn execute_sql_and_expect(instance: &Arc<Instance>, sql: &str, expected: &str) {
1030    let output = execute_sql(instance, sql).await;
1031    let output = output.data.pretty_print().await;
1032    assert_eq!(output, expected.trim());
1033}
tests_integration/test_util.rs

tests_integration/
test_util.rs
