meta_client/client/
procedure.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::future::Future;
16use std::sync::Arc;
17use std::time::Duration;
18
19use api::v1::meta::ddl_task_request::Task;
20use api::v1::meta::procedure_service_client::ProcedureServiceClient;
21use api::v1::meta::{
22    DdlTaskRequest, DdlTaskResponse, GcRegionsRequest, GcRegionsResponse, GcTableRequest,
23    GcTableResponse, MigrateRegionRequest, MigrateRegionResponse, ProcedureActor,
24    ProcedureDetailRequest, ProcedureDetailResponse, ProcedureEventContext, ProcedureId,
25    ProcedureStateResponse, QueryProcedureRequest, ReconcileRequest, ReconcileResponse,
26    RequestHeader, ResponseHeader, Role,
27};
28use common_grpc::channel_manager::ChannelManager;
29use common_meta::procedure_executor::ExecutorContext;
30use common_meta::rpc::ddl::{
31    CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY,
32};
33use common_meta::rpc::procedure::{
34    GcRegionsRequest as MetaGcRegionsRequest, GcResponse as MetaGcResponse,
35    GcTableRequest as MetaGcTableRequest,
36};
37use common_telemetry::tracing_context::TracingContext;
38use common_telemetry::{error, info, warn};
39use snafu::{ResultExt, ensure};
40use tokio::sync::RwLock;
41use tonic::transport::Channel;
42use tonic::{Request, Status};
43
44use crate::client::{Id, LeaderProviderRef, util};
45use crate::error;
46use crate::error::Result;
47
48/// Builds the event context transported by a procedure RPC.
49///
50/// The caller can only supply reason/extensions. Protocol is derived here from
51/// the trusted, typed query channel held locally in the executor context.
52pub(crate) fn procedure_event_context(context: &ExecutorContext) -> Option<ProcedureEventContext> {
53    context.event_input.as_ref().map(|input| {
54        let mut event_context = ProcedureEventContext::from(input);
55        event_context.protocol = context
56            .query_context
57            .as_ref()
58            .and_then(|query_context| query_context.protocol())
59            .unwrap_or_default();
60        event_context
61    })
62}
63
64/// Builds the optional procedure actor transported by a procedure RPC.
65pub(crate) fn procedure_actor(context: &ExecutorContext) -> Option<ProcedureActor> {
66    context.actor.as_ref().map(|username| ProcedureActor {
67        username: username.clone(),
68    })
69}
70
71#[derive(Clone, Debug)]
72pub struct Client {
73    inner: Arc<RwLock<Inner>>,
74}
75
76impl Client {
77    pub fn new(
78        id: Id,
79        role: Role,
80        channel_manager: ChannelManager,
81        max_retry: usize,
82        timeout: Duration,
83    ) -> Self {
84        let inner = Arc::new(RwLock::new(Inner {
85            id,
86            role,
87            channel_manager,
88            leader_provider: None,
89            max_retry,
90            timeout,
91        }));
92
93        Self { inner }
94    }
95
96    /// Start the client with a [LeaderProvider].
97    pub(crate) async fn start_with(&self, leader_provider: LeaderProviderRef) -> Result<()> {
98        let mut inner = self.inner.write().await;
99        inner.start_with(leader_provider)
100    }
101
102    pub async fn submit_ddl_task(&self, req: DdlTaskRequest) -> Result<DdlTaskResponse> {
103        let inner = self.inner.read().await;
104        inner.submit_ddl_task(req).await
105    }
106
107    /// Query the procedure' state by its id
108    pub async fn query_procedure_state(&self, pid: &str) -> Result<ProcedureStateResponse> {
109        let inner = self.inner.read().await;
110        inner.query_procedure_state(pid).await
111    }
112
113    /// Migrate the region from one datanode to the other datanode:
114    /// - `region_id`:  the migrated region id
115    /// - `from_peer`:  the source datanode id
116    /// - `to_peer`:  the target datanode id
117    /// - `timeout`: timeout for downgrading region and upgrading region operations
118    pub async fn migrate_region(
119        &self,
120        context: &ExecutorContext,
121        region_id: u64,
122        from_peer: u64,
123        to_peer: u64,
124        timeout: Duration,
125    ) -> Result<MigrateRegionResponse> {
126        let inner = self.inner.read().await;
127        inner
128            .migrate_region(context, region_id, from_peer, to_peer, timeout)
129            .await
130    }
131
132    /// Reconcile the procedure state.
133    pub async fn reconcile(&self, request: ReconcileRequest) -> Result<ReconcileResponse> {
134        let inner = self.inner.read().await;
135        inner.reconcile(request).await
136    }
137
138    pub async fn list_procedures(&self) -> Result<ProcedureDetailResponse> {
139        let inner = self.inner.read().await;
140        inner.list_procedures().await
141    }
142
143    pub async fn gc_regions(
144        &self,
145        context: &ExecutorContext,
146        request: MetaGcRegionsRequest,
147    ) -> Result<MetaGcResponse> {
148        let inner = self.inner.read().await;
149        inner.gc_regions(context, request).await
150    }
151
152    pub async fn gc_table(
153        &self,
154        context: &ExecutorContext,
155        request: MetaGcTableRequest,
156    ) -> Result<MetaGcResponse> {
157        let inner = self.inner.read().await;
158        inner.gc_table(context, request).await
159    }
160}
161
162#[derive(Debug)]
163struct Inner {
164    id: Id,
165    role: Role,
166    channel_manager: ChannelManager,
167    leader_provider: Option<LeaderProviderRef>,
168    max_retry: usize,
169    /// Request timeout.
170    timeout: Duration,
171}
172
173impl Inner {
174    fn start_with(&mut self, leader_provider: LeaderProviderRef) -> Result<()> {
175        ensure!(
176            !self.is_started(),
177            error::IllegalGrpcClientStateSnafu {
178                err_msg: "DDL client already started",
179            }
180        );
181        self.leader_provider = Some(leader_provider);
182        Ok(())
183    }
184
185    fn make_client(&self, addr: impl AsRef<str>) -> Result<ProcedureServiceClient<Channel>> {
186        let channel = self
187            .channel_manager
188            .get(addr)
189            .context(error::CreateChannelSnafu)?;
190
191        Ok(common_grpc::configure_tonic_client!(
192            ProcedureServiceClient::new(channel),
193            self.channel_manager,
194        ))
195    }
196
197    #[inline]
198    fn is_started(&self) -> bool {
199        self.leader_provider.is_some()
200    }
201
202    async fn with_retry<T, F, R, H>(&self, task: &str, body_fn: F, get_header: H) -> Result<T>
203    where
204        R: Future<Output = std::result::Result<T, Status>>,
205        F: Fn(ProcedureServiceClient<Channel>) -> R,
206        H: Fn(&T) -> &Option<ResponseHeader>,
207    {
208        let Some(leader_provider) = self.leader_provider.as_ref() else {
209            return error::IllegalGrpcClientStateSnafu {
210                err_msg: "not started",
211            }
212            .fail();
213        };
214
215        let mut times = 0;
216        let mut last_error = None;
217
218        while times < self.max_retry {
219            if let Some(leader) = &leader_provider.leader() {
220                let client = self.make_client(leader)?;
221                match body_fn(client).await {
222                    Ok(res) => {
223                        if util::is_not_leader(get_header(&res)) {
224                            last_error = Some(format!("{leader} is not a leader"));
225                            warn!("Failed to {task} to {leader}, not a leader");
226                            let leader = leader_provider.ask_leader().await?;
227                            info!("DDL client updated to new leader addr: {leader}");
228                            times += 1;
229                            continue;
230                        }
231                        return Ok(res);
232                    }
233                    Err(status) => {
234                        // The leader may be unreachable.
235                        if util::is_unreachable(&status) {
236                            last_error = Some(status.to_string());
237                            warn!("Failed to {task} to {leader}, source: {status}");
238                            let leader = leader_provider.ask_leader().await?;
239                            info!("Procedure client updated to new leader addr: {leader}");
240                            times += 1;
241                            continue;
242                        } else {
243                            error!("An error occurred in gRPC, status: {status:?}");
244                            return Err(error::Error::from(status));
245                        }
246                    }
247                }
248            } else {
249                leader_provider.ask_leader().await?;
250            }
251        }
252
253        error::RetryTimesExceededSnafu {
254            msg: format!("Failed to {task}, last error: {:?}", last_error),
255            times: self.max_retry,
256        }
257        .fail()
258    }
259
260    async fn migrate_region(
261        &self,
262        context: &ExecutorContext,
263        region_id: u64,
264        from_peer: u64,
265        to_peer: u64,
266        timeout: Duration,
267    ) -> Result<MigrateRegionResponse> {
268        let mut req = MigrateRegionRequest {
269            region_id,
270            from_peer,
271            to_peer,
272            timeout_secs: timeout.as_secs() as u32,
273            event_context: procedure_event_context(context),
274            actor: procedure_actor(context),
275            ..Default::default()
276        };
277
278        req.set_header(
279            self.id,
280            self.role,
281            TracingContext::from_current_span().to_w3c(),
282        );
283
284        self.with_retry(
285            "migrate region",
286            move |mut client| {
287                let mut req = Request::new(req.clone());
288                req.set_timeout(self.timeout);
289
290                async move { client.migrate(req).await.map(|res| res.into_inner()) }
291            },
292            |resp: &MigrateRegionResponse| &resp.header,
293        )
294        .await
295    }
296
297    async fn reconcile(&self, request: ReconcileRequest) -> Result<ReconcileResponse> {
298        let mut req = request;
299        req.set_header(
300            self.id,
301            self.role,
302            TracingContext::from_current_span().to_w3c(),
303        );
304
305        self.with_retry(
306            "reconcile",
307            move |mut client| {
308                let mut req = Request::new(req.clone());
309                req.set_timeout(self.timeout);
310
311                async move { client.reconcile(req).await.map(|res| res.into_inner()) }
312            },
313            |resp: &ReconcileResponse| &resp.header,
314        )
315        .await
316    }
317
318    async fn gc_regions(
319        &self,
320        context: &ExecutorContext,
321        request: MetaGcRegionsRequest,
322    ) -> Result<MetaGcResponse> {
323        let timeout = request.timeout;
324        let req = GcRegionsRequest {
325            header: Some(RequestHeader {
326                protocol_version: 0,
327                member_id: self.id,
328                role: self.role as i32,
329                tracing_context: TracingContext::from_current_span().to_w3c(),
330            }),
331            region_ids: request.region_ids,
332            full_file_listing: request.full_file_listing,
333            timeout_secs: gc_timeout_secs(timeout),
334            event_context: procedure_event_context(context),
335            actor: procedure_actor(context),
336        };
337
338        let resp: GcRegionsResponse = self
339            .with_retry(
340                "gc_regions",
341                move |mut client| {
342                    let mut req = Request::new(req.clone());
343                    if let Some(timeout) = timeout {
344                        req.set_timeout(timeout);
345                    }
346                    async move { client.gc_regions(req).await.map(|res| res.into_inner()) }
347                },
348                |resp: &GcRegionsResponse| &resp.header,
349            )
350            .await?;
351
352        let stats = resp.stats.unwrap_or_default();
353        Ok(MetaGcResponse {
354            processed_regions: stats.processed_regions,
355            need_retry_regions: stats.need_retry_regions,
356            deleted_files: stats.deleted_files,
357            deleted_indexes: stats.deleted_indexes,
358        })
359    }
360
361    async fn gc_table(
362        &self,
363        context: &ExecutorContext,
364        request: MetaGcTableRequest,
365    ) -> Result<MetaGcResponse> {
366        let timeout = request.timeout;
367        let req = GcTableRequest {
368            header: Some(RequestHeader {
369                protocol_version: 0,
370                member_id: self.id,
371                role: self.role as i32,
372                tracing_context: TracingContext::from_current_span().to_w3c(),
373            }),
374            catalog_name: request.catalog_name,
375            schema_name: request.schema_name,
376            table_name: request.table_name,
377            full_file_listing: request.full_file_listing,
378            timeout_secs: gc_timeout_secs(timeout),
379            event_context: procedure_event_context(context),
380            actor: procedure_actor(context),
381        };
382
383        let resp: GcTableResponse = self
384            .with_retry(
385                "gc_table",
386                move |mut client| {
387                    let mut req = Request::new(req.clone());
388                    if let Some(timeout) = timeout {
389                        req.set_timeout(timeout);
390                    }
391                    async move { client.gc_table(req).await.map(|res| res.into_inner()) }
392                },
393                |resp: &GcTableResponse| &resp.header,
394            )
395            .await?;
396
397        let stats = resp.stats.unwrap_or_default();
398        Ok(MetaGcResponse {
399            processed_regions: stats.processed_regions,
400            need_retry_regions: stats.need_retry_regions,
401            deleted_files: stats.deleted_files,
402            deleted_indexes: stats.deleted_indexes,
403        })
404    }
405
406    async fn query_procedure_state(&self, pid: &str) -> Result<ProcedureStateResponse> {
407        let mut req = QueryProcedureRequest {
408            pid: Some(ProcedureId { key: pid.into() }),
409            ..Default::default()
410        };
411
412        req.set_header(
413            self.id,
414            self.role,
415            TracingContext::from_current_span().to_w3c(),
416        );
417
418        self.with_retry(
419            "query procedure state",
420            move |mut client| {
421                let mut req = Request::new(req.clone());
422                req.set_timeout(self.timeout);
423
424                async move { client.query(req).await.map(|res| res.into_inner()) }
425            },
426            |resp: &ProcedureStateResponse| &resp.header,
427        )
428        .await
429    }
430
431    async fn submit_ddl_task(&self, mut req: DdlTaskRequest) -> Result<DdlTaskResponse> {
432        let creator = create_database_creator_metadata_value(&req);
433        req.set_header(
434            self.id,
435            self.role,
436            TracingContext::from_current_span().to_w3c(),
437        );
438        let timeout = Duration::from_secs(req.timeout_secs.into());
439
440        self.with_retry(
441            "submit ddl task",
442            move |mut client| {
443                let mut req = Request::new(req.clone());
444                if let Some(value) = creator.as_deref() {
445                    req.metadata_mut().insert_bin(
446                        CREATE_DATABASE_CREATOR_METADATA_KEY,
447                        tonic::metadata::MetadataValue::from_bytes(value.as_bytes()),
448                    );
449                }
450                req.set_timeout(timeout);
451                async move { client.ddl(req).await.map(|res| res.into_inner()) }
452            },
453            |resp: &DdlTaskResponse| &resp.header,
454        )
455        .await
456    }
457
458    async fn list_procedures(&self) -> Result<ProcedureDetailResponse> {
459        let mut req = ProcedureDetailRequest::default();
460        req.set_header(
461            self.id,
462            self.role,
463            TracingContext::from_current_span().to_w3c(),
464        );
465
466        self.with_retry(
467            "list procedure",
468            move |mut client| {
469                let mut req = Request::new(req.clone());
470                req.set_timeout(self.timeout);
471                async move { client.details(req).await.map(|res| res.into_inner()) }
472            },
473            |resp: &ProcedureDetailResponse| &resp.header,
474        )
475        .await
476    }
477}
478
479fn create_database_creator_metadata_value(req: &DdlTaskRequest) -> Option<String> {
480    // StatementExecutor removes client values before attaching the authenticated creator.
481    if !matches!(req.task, Some(Task::CreateDatabaseTask(_))) {
482        return None;
483    }
484
485    req.query_context
486        .as_ref()?
487        .extensions
488        .get(CREATE_DATABASE_CREATOR_EXTENSION_KEY)
489        .cloned()
490}
491
492fn gc_timeout_secs(timeout: Option<Duration>) -> u32 {
493    timeout
494        .map(|timeout| timeout.as_secs().max(1).try_into().unwrap_or(u32::MAX))
495        .unwrap_or(0)
496}
497
498#[cfg(test)]
499mod tests {
500    use std::time::{Duration, Instant};
501
502    use api::v1::meta::heartbeat_server::{Heartbeat, HeartbeatServer};
503    use api::v1::meta::procedure_service_server::{ProcedureService, ProcedureServiceServer};
504    use api::v1::meta::{
505        AskLeaderRequest, AskLeaderResponse, DdlTaskRequest, DdlTaskResponse, GcRegionsRequest,
506        GcRegionsResponse, GcTableRequest, GcTableResponse, HeartbeatRequest, HeartbeatResponse,
507        MigrateRegionRequest, MigrateRegionResponse, Peer, ProcedureDetailRequest,
508        ProcedureDetailResponse, ProcedureStateResponse, QueryProcedureRequest, ReconcileRequest,
509        ReconcileResponse, ResponseHeader, Role,
510    };
511    use async_trait::async_trait;
512    use common_base::protocol::Channel;
513    use common_error::status_code::StatusCode;
514    use common_event_recorder::{PersistentEventContext, ProcedureEventInput};
515    use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
516    use common_meta::rpc::ddl::{
517        CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY,
518        CommentObjectType, CommentOnTask, CreatorGrantIntent, DdlTask, QueryContext,
519        SubmitDdlTaskRequest, TriggerReason,
520    };
521    use common_telemetry::common_error::ext::ErrorExt;
522    use common_telemetry::info;
523    use tokio::net::TcpListener;
524    use tokio::sync::mpsc;
525    use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream};
526    use tonic::codec::CompressionEncoding;
527    use tonic::{Request, Response, Status};
528
529    use crate::client::MetaClientBuilder;
530    use crate::client::procedure::{gc_timeout_secs, procedure_event_context};
531
532    #[test]
533    fn test_gc_timeout_secs() {
534        assert_eq!(gc_timeout_secs(None), 0);
535        assert_eq!(gc_timeout_secs(Some(Duration::from_millis(1))), 1);
536        assert_eq!(gc_timeout_secs(Some(Duration::from_millis(999))), 1);
537        assert_eq!(gc_timeout_secs(Some(Duration::from_secs(1))), 1);
538        assert_eq!(gc_timeout_secs(Some(Duration::from_secs(10))), 10);
539    }
540
541    #[test]
542    fn test_procedure_event_context_derives_protocol_from_query_context() {
543        let context = ExecutorContext {
544            query_context: Some(QueryContext {
545                channel: Channel::Postgres as u8,
546                ..Default::default()
547            }),
548            event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)),
549            ..Default::default()
550        };
551
552        assert_eq!(
553            procedure_event_context(&context).map(PersistentEventContext::from),
554            Some(PersistentEventContext::new(TriggerReason::Manual).with_protocol("postgres"))
555        );
556
557        let automatic_context = ExecutorContext {
558            event_input: Some(ProcedureEventInput::new(TriggerReason::ScheduledGc)),
559            ..Default::default()
560        };
561        assert_eq!(
562            procedure_event_context(&automatic_context).map(PersistentEventContext::from),
563            Some(PersistentEventContext::new(TriggerReason::ScheduledGc))
564        );
565
566        let unknown_channel_context = ExecutorContext {
567            query_context: Some(QueryContext::default()),
568            event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)),
569            ..Default::default()
570        };
571        assert_eq!(
572            procedure_event_context(&unknown_channel_context).map(PersistentEventContext::from),
573            Some(PersistentEventContext::new(TriggerReason::Manual))
574        );
575    }
576
577    #[derive(Clone)]
578    struct MockHeartbeat {
579        leader_addr: String,
580    }
581
582    #[async_trait]
583    impl Heartbeat for MockHeartbeat {
584        type HeartbeatStream = ReceiverStream<Result<HeartbeatResponse, Status>>;
585
586        async fn heartbeat(
587            &self,
588            _request: Request<tonic::Streaming<HeartbeatRequest>>,
589        ) -> Result<Response<Self::HeartbeatStream>, Status> {
590            Err(Status::unimplemented(
591                "heartbeat stream is not used in this test",
592            ))
593        }
594
595        async fn ask_leader(
596            &self,
597            _request: Request<AskLeaderRequest>,
598        ) -> Result<Response<AskLeaderResponse>, Status> {
599            Ok(Response::new(AskLeaderResponse {
600                header: Some(ResponseHeader {
601                    protocol_version: 0,
602                    error: None,
603                }),
604                leader: Some(Peer {
605                    id: 1,
606                    addr: self.leader_addr.clone(),
607                }),
608            }))
609        }
610    }
611
612    #[derive(Clone)]
613    struct MockProcedure {
614        delay: Duration,
615        request_tx: Option<mpsc::UnboundedSender<Request<DdlTaskRequest>>>,
616    }
617
618    #[async_trait]
619    impl ProcedureService for MockProcedure {
620        async fn query(
621            &self,
622            _request: Request<QueryProcedureRequest>,
623        ) -> Result<Response<ProcedureStateResponse>, Status> {
624            Err(Status::unimplemented("query is not used in this test"))
625        }
626
627        async fn ddl(
628            &self,
629            request: Request<DdlTaskRequest>,
630        ) -> Result<Response<DdlTaskResponse>, Status> {
631            if let Some(request_tx) = &self.request_tx {
632                request_tx.send(request).unwrap();
633            }
634            tokio::time::sleep(self.delay).await;
635            Ok(Response::new(DdlTaskResponse {
636                header: Some(ResponseHeader {
637                    protocol_version: 0,
638                    error: None,
639                }),
640                ..Default::default()
641            }))
642        }
643
644        async fn reconcile(
645            &self,
646            _request: Request<ReconcileRequest>,
647        ) -> Result<Response<ReconcileResponse>, Status> {
648            Err(Status::unimplemented("reconcile is not used in this test"))
649        }
650
651        async fn migrate(
652            &self,
653            _request: Request<MigrateRegionRequest>,
654        ) -> Result<Response<MigrateRegionResponse>, Status> {
655            Err(Status::unimplemented("migrate is not used in this test"))
656        }
657
658        async fn details(
659            &self,
660            _request: Request<ProcedureDetailRequest>,
661        ) -> Result<Response<ProcedureDetailResponse>, Status> {
662            Err(Status::unimplemented("details is not used in this test"))
663        }
664
665        async fn gc_regions(
666            &self,
667            _request: Request<GcRegionsRequest>,
668        ) -> Result<Response<GcRegionsResponse>, Status> {
669            Err(Status::unimplemented("gc_regions is not used in this test"))
670        }
671
672        async fn gc_table(
673            &self,
674            _request: Request<GcTableRequest>,
675        ) -> Result<Response<GcTableResponse>, Status> {
676            Err(Status::unimplemented("gc_table is not used in this test"))
677        }
678    }
679
680    #[tokio::test(flavor = "multi_thread")]
681    async fn test_meta_client_forwards_create_database_creator_metadata() {
682        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
683        let addr_str = listener.local_addr().unwrap().to_string();
684        let (request_tx, mut request_rx) = mpsc::unbounded_channel();
685        let heartbeat = MockHeartbeat {
686            leader_addr: addr_str.clone(),
687        };
688        let procedure = MockProcedure {
689            delay: Duration::ZERO,
690            request_tx: Some(request_tx),
691        };
692        let server = tonic::transport::Server::builder()
693            .add_service(
694                HeartbeatServer::new(heartbeat).accept_compressed(CompressionEncoding::Zstd),
695            )
696            .add_service(
697                ProcedureServiceServer::new(procedure).accept_compressed(CompressionEncoding::Zstd),
698            )
699            .serve_with_incoming(TcpListenerStream::new(listener));
700        let server_handle = tokio::spawn(server);
701
702        let mut client = MetaClientBuilder::new(0, Role::Frontend)
703            .enable_heartbeat()
704            .enable_procedure()
705            .build();
706        client.start(&[addr_str.as_str()]).await.unwrap();
707
708        let creator = CreatorGrantIntent {
709            username: "alice".to_string(),
710            created_at_ns: 42,
711        };
712        let executor_context = ExecutorContext {
713            query_context: Some(QueryContext {
714                channel: Channel::Postgres as u8,
715                ..Default::default()
716            }),
717            actor: Some("effective-user".to_string()),
718            event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)),
719            ..Default::default()
720        };
721        ProcedureExecutor::submit_ddl_task(
722            &client,
723            executor_context,
724            SubmitDdlTaskRequest::new(DdlTask::new_create_database(
725                "greptime".to_string(),
726                "metrics".to_string(),
727                false,
728                Default::default(),
729                Some(creator.clone()),
730            )),
731        )
732        .await
733        .unwrap();
734
735        let request = request_rx.recv().await.unwrap();
736        let encoded = serde_json::to_string(&creator).unwrap();
737        assert_eq!(
738            request
739                .metadata()
740                .get_bin(CREATE_DATABASE_CREATOR_METADATA_KEY)
741                .unwrap()
742                .to_bytes()
743                .unwrap()
744                .as_ref(),
745            encoded.as_bytes()
746        );
747        let request = request.into_inner();
748        assert_eq!(request.actor.unwrap().username, "effective-user");
749        assert_eq!(
750            PersistentEventContext::from(request.event_context.unwrap()),
751            PersistentEventContext::new(TriggerReason::Manual).with_protocol("postgres")
752        );
753        let extensions = &request.query_context.unwrap().extensions;
754        assert_eq!(extensions[CREATE_DATABASE_CREATOR_EXTENSION_KEY], encoded);
755
756        server_handle.abort();
757    }
758
759    #[tokio::test(flavor = "multi_thread")]
760    async fn test_meta_client_ddl_request_timeout() {
761        common_telemetry::init_default_ut_logging();
762
763        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
764        let addr = listener.local_addr().unwrap();
765        let addr_str = addr.to_string();
766
767        let heartbeat = MockHeartbeat {
768            leader_addr: addr_str.clone(),
769        };
770        let procedure = MockProcedure {
771            delay: Duration::from_secs(4),
772            request_tx: None,
773        };
774
775        let server = tonic::transport::Server::builder()
776            .add_service(
777                HeartbeatServer::new(heartbeat)
778                    .accept_compressed(CompressionEncoding::Gzip)
779                    .accept_compressed(CompressionEncoding::Zstd),
780            )
781            .add_service(
782                ProcedureServiceServer::new(procedure)
783                    .accept_compressed(CompressionEncoding::Gzip)
784                    .accept_compressed(CompressionEncoding::Zstd),
785            )
786            .serve_with_incoming(TcpListenerStream::new(listener));
787        let server_handle = tokio::spawn(server);
788
789        let mut client = MetaClientBuilder::new(0, Role::Frontend)
790            .enable_heartbeat()
791            .enable_procedure()
792            .build();
793        client.start(&[addr_str.as_str()]).await.unwrap();
794
795        let mut request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(CommentOnTask {
796            catalog_name: "greptime".to_string(),
797            schema_name: "public".to_string(),
798            object_type: CommentObjectType::Table,
799            object_name: "test_table".to_string(),
800            column_name: None,
801            object_id: None,
802            comment: Some("timeout".to_string()),
803        }));
804        request.timeout = Duration::from_secs(1);
805
806        let now = Instant::now();
807        let err = client
808            .submit_ddl_task(
809                ExecutorContext {
810                    query_context: Some(QueryContext::default()),
811                    ..Default::default()
812                },
813                request,
814            )
815            .await
816            .unwrap_err();
817        let elapsed = now.elapsed();
818        // The request should be cancelled within 1 second.
819        assert!(elapsed < Duration::from_secs(2));
820        info!("err: {err:?}, code: {}", err.status_code());
821        assert_eq!(err.status_code(), StatusCode::Cancelled);
822        let err_msg = err.to_string();
823        assert!(
824            err_msg.contains("Timeout expired"),
825            "unexpected error: {err_msg}"
826        );
827
828        server_handle.abort();
829    }
830}
meta_client/client/procedure.rs

meta_client/client/
procedure.rs
