frontend/instance/
region_query.rs1use 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}