Skip to main content

frontend/instance/
region_query.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::sync::Arc;
16
17use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
18use async_trait::async_trait;
19use client::region::{
20    build_remote_dyn_filter_unregister_request, build_remote_dyn_filter_update_request,
21};
22use common_error::ext::BoxedError;
23use common_meta::node_manager::NodeManagerRef;
24use common_query::request::QueryRequest;
25use common_recordbatch::SendableRecordBatchStream;
26use partition::manager::PartitionRuleManagerRef;
27use query::error::{RegionQuerySnafu, Result as QueryResult};
28use query::region_query::RegionQueryHandler;
29use session::ReadPreference;
30use snafu::ResultExt;
31use store_api::storage::RegionId;
32
33use crate::error::{FindRegionPeerSnafu, RequestQuerySnafu, Result};
34
35pub(crate) struct FrontendRegionQueryHandler {
36    partition_manager: PartitionRuleManagerRef,
37    node_manager: NodeManagerRef,
38}
39
40impl FrontendRegionQueryHandler {
41    pub fn arc(
42        partition_manager: PartitionRuleManagerRef,
43        node_manager: NodeManagerRef,
44    ) -> Arc<Self> {
45        Arc::new(Self {
46            partition_manager,
47            node_manager,
48        })
49    }
50}
51
52#[async_trait]
53impl RegionQueryHandler for FrontendRegionQueryHandler {
54    async fn do_get(
55        &self,
56        read_preference: ReadPreference,
57        request: QueryRequest,
58    ) -> QueryResult<SendableRecordBatchStream> {
59        self.do_get_inner(read_preference, request)
60            .await
61            .map_err(BoxedError::new)
62            .context(RegionQuerySnafu)
63    }
64
65    async fn handle_remote_dyn_filter_update(
66        &self,
67        region_id: RegionId,
68        query_id: String,
69        update: RemoteDynFilterUpdate,
70    ) -> QueryResult<()> {
71        self.handle_remote_dyn_filter_update_inner(region_id, query_id, update)
72            .await
73            .map_err(BoxedError::new)
74            .context(RegionQuerySnafu)
75    }
76
77    async fn handle_remote_dyn_filter_unregister(
78        &self,
79        region_id: RegionId,
80        query_id: String,
81        unregister: RemoteDynFilterUnregister,
82    ) -> QueryResult<()> {
83        self.handle_remote_dyn_filter_unregister_inner(region_id, query_id, unregister)
84            .await
85            .map_err(BoxedError::new)
86            .context(RegionQuerySnafu)
87    }
88}
89
90impl FrontendRegionQueryHandler {
91    async fn do_get_inner(
92        &self,
93        read_preference: ReadPreference,
94        request: QueryRequest,
95    ) -> Result<SendableRecordBatchStream> {
96        let region_id = request.region_id;
97
98        let peer = &self
99            .partition_manager
100            .find_region_leader(region_id)
101            .await
102            .context(FindRegionPeerSnafu {
103                region_id,
104                read_preference,
105            })?;
106
107        let client = self.node_manager.datanode(peer).await;
108
109        client
110            .handle_query(request)
111            .await
112            .context(RequestQuerySnafu)
113    }
114
115    async fn handle_remote_dyn_filter_update_inner(
116        &self,
117        region_id: RegionId,
118        query_id: String,
119        update: RemoteDynFilterUpdate,
120    ) -> Result<()> {
121        let peer = &self
122            .partition_manager
123            .find_region_leader(region_id)
124            .await
125            .context(FindRegionPeerSnafu {
126                region_id,
127                read_preference: ReadPreference::Leader,
128            })?;
129        let client = self.node_manager.datanode(peer).await;
130        client
131            .handle(build_remote_dyn_filter_update_request(query_id, update))
132            .await
133            .context(RequestQuerySnafu)?;
134        Ok(())
135    }
136
137    async fn handle_remote_dyn_filter_unregister_inner(
138        &self,
139        region_id: RegionId,
140        query_id: String,
141        unregister: RemoteDynFilterUnregister,
142    ) -> Result<()> {
143        let peer = &self
144            .partition_manager
145            .find_region_leader(region_id)
146            .await
147            .context(FindRegionPeerSnafu {
148                region_id,
149                read_preference: ReadPreference::Leader,
150            })?;
151        let client = self.node_manager.datanode(peer).await;
152        client
153            .handle(build_remote_dyn_filter_unregister_request(
154                query_id, unregister,
155            ))
156            .await
157            .context(RequestQuerySnafu)?;
158        Ok(())
159    }
160}