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            providers,
366            store: default_store,
367        },
368        grpc: GrpcOptions::default()
369            .with_bind_addr(PEER_PLACEHOLDER_ADDR)
370            .with_server_addr(PEER_PLACEHOLDER_ADDR),
371        wal: wal_config,
372        region_engine,
373        ..Default::default()
374    }
375}
376
377pub(crate) async fn create_test_table(instance: &Instance, table_name: &str) {
378    let sql = format!(
379        r#"
380CREATE TABLE IF NOT EXISTS {table_name} (
381    host String NOT NULL PRIMARY KEY,
382    cpu DOUBLE NULL,
383    memory DOUBLE NULL,
384    ts TIMESTAMP NOT NULL TIME INDEX,
385)
386"#
387    );
388
389    let result = instance.do_query(&sql, QueryContext::arc()).await;
390    let _ = result.first().unwrap().as_ref().unwrap();
391}
392
393async fn setup_standalone_instance(
394    test_name: &str,
395    store_type: StorageType,
396) -> GreptimeDbStandalone {
397    GreptimeDbStandaloneBuilder::new(test_name)
398        .with_default_store_type(store_type)
399        .build()
400        .await
401}
402
403async fn setup_standalone_instance_with_slow_query_threshold(
404    test_name: &str,
405    store_type: StorageType,
406    slow_query_threshold: std::time::Duration,
407) -> GreptimeDbStandalone {
408    GreptimeDbStandaloneBuilder::new(test_name)
409        .with_default_store_type(store_type)
410        .with_slow_query_threshold(slow_query_threshold)
411        .build()
412        .await
413}
414
415async fn setup_standalone_instance_with_plugins(
416    test_name: &str,
417    store_type: StorageType,
418    plugins: Plugins,
419) -> GreptimeDbStandalone {
420    GreptimeDbStandaloneBuilder::new(test_name)
421        .with_default_store_type(store_type)
422        .with_plugin(plugins)
423        .build()
424        .await
425}
426
427async fn setup_standalone_instance_with_plugins_and_slow_query_threshold(
428    test_name: &str,
429    store_type: StorageType,
430    plugins: Plugins,
431    slow_query_threshold: std::time::Duration,
432) -> GreptimeDbStandalone {
433    GreptimeDbStandaloneBuilder::new(test_name)
434        .with_default_store_type(store_type)
435        .with_plugin(plugins)
436        .with_slow_query_threshold(slow_query_threshold)
437        .build()
438        .await
439}
440
441pub async fn setup_test_http_app(store_type: StorageType, name: &str) -> (Router, TestGuard) {
442    let instance = setup_standalone_instance(name, store_type).await;
443
444    let http_opts = HttpOptions {
445        addr: format!("127.0.0.1:{}", ports::get_port()),
446        ..Default::default()
447    };
448    let http_server = HttpServerBuilder::new(http_opts)
449        .with_sql_handler(instance.fe_instance().clone())
450        .with_logs_handler(instance.fe_instance().clone())
451        .with_metrics_handler(MetricsHandler)
452        .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap())
453        .build();
454    (
455        http_server.build(http_server.make_app()).unwrap(),
456        instance.guard,
457    )
458}
459
460pub async fn setup_test_http_app_with_frontend(
461    store_type: StorageType,
462    name: &str,
463) -> (Router, TestGuard) {
464    setup_test_http_app_with_frontend_and_user_provider(store_type, name, None).await
465}
466
467pub async fn setup_test_http_app_with_frontend_and_slow_query_threshold(
468    store_type: StorageType,
469    name: &str,
470    slow_query_threshold: std::time::Duration,
471) -> (Router, TestGuard) {
472    let instance =
473        setup_standalone_instance_with_slow_query_threshold(name, store_type, slow_query_threshold)
474            .await;
475
476    create_test_table(instance.fe_instance(), "demo").await;
477
478    let http_opts = HttpOptions {
479        addr: format!("127.0.0.1:{}", ports::get_port()),
480        ..Default::default()
481    };
482
483    let http_server = HttpServerBuilder::new(http_opts)
484        .with_sql_handler(instance.fe_instance().clone())
485        .with_log_ingest_handler(instance.fe_instance().clone(), None, None)
486        .with_logs_handler(instance.fe_instance().clone())
487        .with_influxdb_handler(instance.fe_instance().clone())
488        .with_otlp_handler(instance.fe_instance().clone(), true)
489        .with_jaeger_handler(instance.fe_instance().clone())
490        .with_greptime_config_options(instance.opts.to_toml().unwrap())
491        .build();
492
493    let app = http_server.build(http_server.make_app()).unwrap();
494    (app, instance.guard)
495}
496
497pub async fn setup_test_http_app_with_frontend_and_user_provider(
498    store_type: StorageType,
499    name: &str,
500    user_provider: Option<UserProviderRef>,
501) -> (Router, TestGuard) {
502    setup_test_http_app_with_frontend_and_custom_options(
503        store_type,
504        name,
505        user_provider,
506        None,
507        None,
508    )
509    .await
510}
511
512pub async fn setup_test_http_app_with_frontend_and_custom_options(
513    store_type: StorageType,
514    name: &str,
515    user_provider: Option<UserProviderRef>,
516    http_opts: Option<HttpOptions>,
517    memory_limiter: Option<ServerMemoryLimiter>,
518) -> (Router, TestGuard) {
519    let plugins = Plugins::new();
520    if let Some(user_provider) = user_provider.clone() {
521        plugins.insert::<UserProviderRef>(user_provider.clone());
522        plugins.insert::<PermissionCheckerRef>(DefaultPermissionChecker::arc());
523    }
524
525    let instance = setup_standalone_instance_with_plugins(name, store_type, plugins).await;
526
527    create_test_table(instance.fe_instance(), "demo").await;
528
529    let http_opts = http_opts.unwrap_or_else(|| HttpOptions {
530        addr: format!("127.0.0.1:{}", ports::get_port()),
531        ..Default::default()
532    });
533
534    let mut http_server = HttpServerBuilder::new(http_opts)
535        .with_sql_handler(instance.fe_instance().clone())
536        .with_log_ingest_handler(instance.fe_instance().clone(), None, None)
537        .with_logs_handler(instance.fe_instance().clone())
538        .with_influxdb_handler(instance.fe_instance().clone())
539        .with_otlp_handler(instance.fe_instance().clone(), true)
540        .with_jaeger_handler(instance.fe_instance().clone())
541        .with_dashboard_handler(instance.fe_instance().clone())
542        .with_greptime_config_options(instance.opts.to_toml().unwrap());
543
544    if let Some(user_provider) = user_provider {
545        http_server = http_server.with_user_provider(user_provider);
546    }
547
548    if let Some(limiter) = memory_limiter {
549        http_server = http_server.with_memory_limiter(limiter);
550    }
551
552    let http_server = http_server.build();
553
554    let app = http_server.build(http_server.make_app()).unwrap();
555    (app, instance.guard)
556}
557
558async fn run_sql(sql: &str, instance: &GreptimeDbStandalone) {
559    let result = instance
560        .fe_instance()
561        .do_query(sql, QueryContext::arc())
562        .await;
563    let _ = result.first().unwrap().as_ref().unwrap();
564}
565
566pub async fn setup_test_prom_app_with_frontend(
567    store_type: StorageType,
568    name: &str,
569) -> (Router, TestGuard) {
570    setup_test_prom_app_with_frontend_inner(store_type, name, false).await
571}
572
573/// Like [`setup_test_prom_app_with_frontend`] but enables the pending-rows batcher,
574/// so Prometheus remote write goes through the batched (metric-engine) path instead
575/// of the direct `PromStoreProtocolHandler::write` path.
576pub async fn setup_test_prom_app_with_frontend_batched(
577    store_type: StorageType,
578    name: &str,
579) -> (Router, TestGuard) {
580    setup_test_prom_app_with_frontend_inner(store_type, name, true).await
581}
582
583async fn setup_test_prom_app_with_frontend_inner(
584    store_type: StorageType,
585    name: &str,
586    enable_batcher: bool,
587) -> (Router, TestGuard) {
588    unsafe {
589        std::env::set_var("TZ", "UTC");
590    }
591
592    let instance = setup_standalone_instance(name, store_type).await;
593
594    // build physical table
595    let sql = "CREATE TABLE phy (ts timestamp time index, val double, host string primary key) engine=metric with ('physical_metric_table' = '')";
596    run_sql(sql, &instance).await;
597    let sql = "CREATE TABLE phy_ns (ts timestamp(0) time index, val double, host string primary key) engine=metric with ('physical_metric_table' = '')";
598    run_sql(sql, &instance).await;
599    // build metric tables
600    let sql = "CREATE TABLE demo (ts timestamp time index, val double, host string primary key) engine=metric with ('on_physical_table' = 'phy')";
601    run_sql(sql, &instance).await;
602    let sql = "CREATE TABLE demo_metrics (ts timestamp time index, val double, idc string primary key) engine=metric with ('on_physical_table' = 'phy')";
603    run_sql(sql, &instance).await;
604    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')";
605    run_sql(sql, &instance).await;
606
607    // insert rows
608    let sql = "INSERT INTO demo(host, val, ts) VALUES ('host1', 1.1, 0), ('host2', 2.1, 600000)";
609    run_sql(sql, &instance).await;
610    let sql =
611        "INSERT INTO demo_metrics(idc, val, ts) VALUES ('idc1', 1.1, 0), ('idc2', 2.1, 600000)";
612    run_sql(sql, &instance).await;
613    // insert a row with empty label
614    let sql = "INSERT INTO demo_metrics(val, ts) VALUES (1.1, 0)";
615    run_sql(sql, &instance).await;
616
617    // insert rows to multi_labels
618    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)";
619    run_sql(sql, &instance).await;
620
621    // build physical table
622    let sql = "CREATE TABLE phy2 (ts timestamp(9) time index, val double, host string primary key) engine=metric with ('physical_metric_table' = '')";
623    run_sql(sql, &instance).await;
624    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')";
625    run_sql(sql, &instance).await;
626    let sql = "INSERT INTO demo_metrics_with_nanos(idc, val, ts) VALUES ('idc1', 1.1, 0)";
627    run_sql(sql, &instance).await;
628
629    // a mito table with non-prometheus compatible values
630    let sql = "CREATE TABLE mito (ts timestamp(9) time index, val double, host bigint primary key) engine=mito";
631    run_sql(sql, &instance).await;
632    let sql = "INSERT INTO mito(host, val, ts) VALUES (1, 1.1, 0)";
633    run_sql(sql, &instance).await;
634
635    let http_opts = HttpOptions {
636        addr: format!("127.0.0.1:{}", ports::get_port()),
637        ..Default::default()
638    };
639    let frontend_ref = instance.fe_instance().clone();
640    // Mirror the production wiring at `frontend::server`: build the batcher from the
641    // instance's managers. A short flush interval keeps the test responsive.
642    let pending_rows_batcher = if enable_batcher {
643        PendingRowsBatcher::try_new(
644            frontend_ref.partition_manager().clone(),
645            frontend_ref.node_manager().clone(),
646            frontend_ref.catalog_manager().clone(),
647            true,
648            frontend_ref.clone(),
649            Duration::from_millis(50),
650            1000,
651            4,
652            64,
653            64,
654        )
655    } else {
656        None
657    };
658    let http_server = HttpServerBuilder::new(http_opts)
659        .with_sql_handler(frontend_ref.clone())
660        .with_logs_handler(instance.fe_instance().clone())
661        .with_prom_handler(
662            frontend_ref.clone(),
663            Some(frontend_ref.clone()),
664            true,
665            PromValidationMode::Strict,
666            pending_rows_batcher,
667        )
668        .with_prometheus_handler(frontend_ref)
669        .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap())
670        .build();
671    let app = http_server.build(http_server.make_app()).unwrap();
672    (app, instance.guard)
673}
674
675pub async fn setup_grpc_server(
676    store_type: StorageType,
677    name: &str,
678) -> (GreptimeDbStandalone, Arc<GrpcServer>) {
679    setup_grpc_server_with(store_type, name, None, None, None).await
680}
681
682pub async fn setup_grpc_server_with_user_provider(
683    store_type: StorageType,
684    name: &str,
685    user_provider: Option<UserProviderRef>,
686) -> (GreptimeDbStandalone, Arc<GrpcServer>) {
687    setup_grpc_server_with(store_type, name, user_provider, None, None).await
688}
689
690/// Sets up a gRPC server backed by a standalone instance whose frontend has auto
691/// table creation disabled, for testing the server-side global switch.
692pub async fn setup_grpc_server_with_auto_create_table_disabled(
693    store_type: StorageType,
694    name: &str,
695) -> (GreptimeDbStandalone, Arc<GrpcServer>) {
696    let instance = GreptimeDbStandaloneBuilder::new(name)
697        .with_default_store_type(store_type)
698        .with_auto_create_table(false)
699        .build()
700        .await;
701    setup_grpc_server_for_instance(instance, None, None, None).await
702}
703
704pub async fn setup_grpc_server_with(
705    store_type: StorageType,
706    name: &str,
707    user_provider: Option<UserProviderRef>,
708    grpc_config: Option<GrpcServerConfig>,
709    memory_limiter: Option<servers::request_memory_limiter::ServerMemoryLimiter>,
710) -> (GreptimeDbStandalone, Arc<GrpcServer>) {
711    let instance = setup_standalone_instance(name, store_type).await;
712    setup_grpc_server_for_instance(instance, user_provider, grpc_config, memory_limiter).await
713}
714
715/// Builds and starts a gRPC server on top of an already-constructed standalone
716/// instance. This is the shared core behind the `setup_grpc_server_*` helpers.
717async fn setup_grpc_server_for_instance(
718    instance: GreptimeDbStandalone,
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 runtime: Runtime = RuntimeBuilder::default()
724        .worker_threads(2)
725        .thread_name("grpc-handlers")
726        .build()
727        .unwrap();
728
729    let fe_instance_ref = instance.fe_instance().clone();
730
731    let greptime_request_handler = GreptimeRequestHandler::new(
732        fe_instance_ref.clone(),
733        user_provider.clone(),
734        Some(runtime.clone()),
735        FlightCompression::default(),
736    );
737
738    let flight_handler = Arc::new(greptime_request_handler.clone());
739
740    let grpc_config = grpc_config.unwrap_or_default();
741    let mut grpc_builder = GrpcServerBuilder::new(grpc_config.clone(), runtime);
742
743    if let Some(limiter) = memory_limiter {
744        grpc_builder = grpc_builder.with_memory_limiter(limiter);
745    }
746
747    let grpc_builder = grpc_builder
748        .database_handler(greptime_request_handler)
749        .flight_handler(flight_handler)
750        .prometheus_handler(fe_instance_ref.clone(), user_provider.clone())
751        .otel_arrow_handler(OtelArrowServiceHandler::new(fe_instance_ref, user_provider))
752        .with_tls_config(grpc_config.tls)
753        .unwrap();
754
755    let mut grpc_server = grpc_builder.build();
756
757    let fe_grpc_addr = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
758    grpc_server.start(fe_grpc_addr).await.unwrap();
759
760    (instance, Arc::new(grpc_server))
761}
762
763pub async fn setup_mysql_server(
764    store_type: StorageType,
765    name: &str,
766) -> (TestGuard, Arc<Box<dyn Server>>) {
767    setup_mysql_server_with_user_provider(store_type, name, None).await
768}
769
770pub async fn setup_mysql_server_with_slow_query_threshold(
771    store_type: StorageType,
772    name: &str,
773    slow_query_threshold: std::time::Duration,
774) -> (TestGuard, Arc<Box<dyn Server>>) {
775    let plugins = Plugins::new();
776    let instance = setup_standalone_instance_with_plugins_and_slow_query_threshold(
777        name,
778        store_type,
779        plugins,
780        slow_query_threshold,
781    )
782    .await;
783
784    let runtime = RuntimeBuilder::default()
785        .worker_threads(2)
786        .thread_name("mysql-runtime")
787        .build()
788        .unwrap();
789
790    let fe_mysql_addr = format!("127.0.0.1:{}", ports::get_port());
791
792    let fe_instance_ref = instance.fe_instance().clone();
793    let opts = MysqlOptions {
794        addr: fe_mysql_addr.clone(),
795        ..Default::default()
796    };
797    let mut mysql_server = MysqlServer::create_server(
798        runtime,
799        Arc::new(MysqlSpawnRef::new(fe_instance_ref, None)),
800        Arc::new(MysqlSpawnConfig::new(
801            false,
802            Arc::new(
803                ReloadableTlsServerConfig::try_new(opts.tls.clone())
804                    .expect("Failed to load certificates and keys"),
805            ),
806            0,
807            opts.reject_no_database.unwrap_or(false),
808            opts.prepared_stmt_cache_size,
809        )),
810        None,
811    );
812
813    mysql_server
814        .start(fe_mysql_addr.parse::<SocketAddr>().unwrap())
815        .await
816        .unwrap();
817
818    (instance.guard, Arc::new(mysql_server))
819}
820
821pub async fn setup_mysql_server_with_user_provider(
822    store_type: StorageType,
823    name: &str,
824    user_provider: Option<UserProviderRef>,
825) -> (TestGuard, Arc<Box<dyn Server>>) {
826    let plugins = Plugins::new();
827    if let Some(user_provider) = user_provider.clone() {
828        plugins.insert::<UserProviderRef>(user_provider.clone());
829        plugins.insert::<PermissionCheckerRef>(DefaultPermissionChecker::arc());
830    }
831
832    let instance = setup_standalone_instance_with_plugins(name, store_type, plugins).await;
833
834    let runtime = RuntimeBuilder::default()
835        .worker_threads(2)
836        .thread_name("mysql-runtime")
837        .build()
838        .unwrap();
839
840    let fe_mysql_addr = format!("127.0.0.1:{}", ports::get_port());
841
842    let fe_instance_ref = instance.fe_instance().clone();
843    let opts = MysqlOptions {
844        addr: fe_mysql_addr.clone(),
845        ..Default::default()
846    };
847    let mut mysql_server = MysqlServer::create_server(
848        runtime,
849        Arc::new(MysqlSpawnRef::new(fe_instance_ref, user_provider)),
850        Arc::new(MysqlSpawnConfig::new(
851            false,
852            Arc::new(
853                ReloadableTlsServerConfig::try_new(opts.tls.clone())
854                    .expect("Failed to load certificates and keys"),
855            ),
856            0,
857            opts.reject_no_database.unwrap_or(false),
858            opts.prepared_stmt_cache_size,
859        )),
860        None,
861    );
862
863    mysql_server
864        .start(fe_mysql_addr.parse::<SocketAddr>().unwrap())
865        .await
866        .unwrap();
867
868    (instance.guard, Arc::new(mysql_server))
869}
870
871pub async fn setup_pg_server(
872    store_type: StorageType,
873    name: &str,
874) -> (TestGuard, Arc<Box<dyn Server>>) {
875    setup_pg_server_with_user_provider(store_type, name, None).await
876}
877
878pub async fn setup_pg_server_with_slow_query_threshold(
879    store_type: StorageType,
880    name: &str,
881    slow_query_threshold: std::time::Duration,
882) -> (TestGuard, Arc<Box<dyn Server>>) {
883    let instance =
884        setup_standalone_instance_with_slow_query_threshold(name, store_type, slow_query_threshold)
885            .await;
886
887    let runtime = RuntimeBuilder::default()
888        .worker_threads(2)
889        .thread_name("pg-runtime")
890        .build()
891        .unwrap();
892
893    let fe_pg_addr = format!("127.0.0.1:{}", ports::get_port());
894
895    let fe_instance_ref = instance.fe_instance().clone();
896    let opts = PostgresOptions {
897        addr: fe_pg_addr.clone(),
898        ..Default::default()
899    };
900    let tls_server_config = Arc::new(
901        ReloadableTlsServerConfig::try_new(opts.tls.clone())
902            .expect("Failed to load certificates and keys"),
903    );
904
905    let mut pg_server = Box::new(PostgresServer::new(
906        fe_instance_ref,
907        opts.tls.should_force_tls(),
908        tls_server_config,
909        0,
910        runtime,
911        None,
912        None,
913    ));
914
915    pg_server
916        .start(fe_pg_addr.parse::<SocketAddr>().unwrap())
917        .await
918        .unwrap();
919
920    (instance.guard, Arc::new(pg_server))
921}
922
923pub async fn setup_pg_server_with_user_provider(
924    store_type: StorageType,
925    name: &str,
926    user_provider: Option<UserProviderRef>,
927) -> (TestGuard, Arc<Box<dyn Server>>) {
928    let instance = setup_standalone_instance(name, store_type).await;
929
930    let runtime = RuntimeBuilder::default()
931        .worker_threads(2)
932        .thread_name("pg-runtime")
933        .build()
934        .unwrap();
935
936    let fe_pg_addr = format!("127.0.0.1:{}", ports::get_port());
937
938    let fe_instance_ref = instance.fe_instance().clone();
939    let opts = PostgresOptions {
940        addr: fe_pg_addr.clone(),
941        ..Default::default()
942    };
943    let tls_server_config = Arc::new(
944        ReloadableTlsServerConfig::try_new(opts.tls.clone())
945            .expect("Failed to load certificates and keys"),
946    );
947
948    let mut pg_server = Box::new(PostgresServer::new(
949        fe_instance_ref,
950        opts.tls.should_force_tls(),
951        tls_server_config,
952        0,
953        runtime,
954        user_provider,
955        None,
956    ));
957
958    pg_server
959        .start(fe_pg_addr.parse::<SocketAddr>().unwrap())
960        .await
961        .unwrap();
962
963    (instance.guard, Arc::new(pg_server))
964}
965
966pub(crate) async fn prepare_another_catalog_and_schema(instance: &Instance) {
967    let catalog_manager = instance
968        .catalog_manager()
969        .as_any()
970        .downcast_ref::<KvBackendCatalogManager>()
971        .unwrap();
972
973    let table_metadata_manager = catalog_manager.table_metadata_manager_ref();
974    prepare_another_catalog_and_schema_with_manager(table_metadata_manager).await;
975}
976
977pub(crate) async fn prepare_another_catalog_and_schema_with_kv_backend(kv_backend: KvBackendRef) {
978    let table_metadata_manager = TableMetadataManager::new(kv_backend);
979    prepare_another_catalog_and_schema_with_manager(&table_metadata_manager).await;
980}
981
982async fn prepare_another_catalog_and_schema_with_manager(
983    table_metadata_manager: &TableMetadataManager,
984) {
985    table_metadata_manager
986        .catalog_manager()
987        .create(CatalogNameKey::new("another_catalog"), true)
988        .await
989        .unwrap();
990    table_metadata_manager
991        .schema_manager()
992        .create(
993            SchemaNameKey::new("another_catalog", "another_schema"),
994            None,
995            true,
996        )
997        .await
998        .unwrap();
999}
1000
1001pub async fn execute_sql(instance: &Arc<Instance>, sql: &str) -> Output {
1002    SqlQueryHandler::do_query(instance.as_ref(), sql, QueryContext::arc())
1003        .await
1004        .remove(0)
1005        .unwrap()
1006}
1007
1008pub async fn try_execute_sql(
1009    instance: &Arc<Instance>,
1010    sql: &str,
1011) -> servers::error::Result<Output> {
1012    SqlQueryHandler::do_query(instance.as_ref(), sql, QueryContext::arc())
1013        .await
1014        .remove(0)
1015}
1016
1017pub async fn execute_sql_and_expect(instance: &Arc<Instance>, sql: &str, expected: &str) {
1018    let output = execute_sql(instance, sql).await;
1019    let output = output.data.pretty_print().await;
1020    assert_eq!(output, expected.trim());
1021}
tests_integration/test_util.rs

tests_integration/
test_util.rs
