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

tests_integration/
test_util.rs
