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

tests_integration/
test_util.rs
