tin_sea_geo/lib.rs
1use sea_orm::{EntityTrait, QueryFilter, Select, prelude::Expr};
2
3pub use tin_sea_geo_macros::{with_centroid, with_geo, with_geometry};
4
5/// 空间查询扩展 trait
6///
7/// 为 SeaORM 的 Select 查询添加地理空间查询功能
8pub trait SpatialQueryExt<E>
9where
10 E: EntityTrait,
11{
12 /// 点包含查询 - 查找包含指定点的几何对象
13 ///
14 /// 使用默认的 `geom` 字段和 WGS84(4326) 到 CGCS2000(4527) 坐标转换
15 ///
16 /// # 参数
17 /// * `lng` - 经度
18 /// * `lat` - 纬度
19 ///
20 /// # 示例
21 /// ```
22 /// use my_macros::SpatialQueryExt;
23 ///
24 /// let result = Entity::find()
25 /// .contains(116.26953, 39.74616)
26 /// .all(&db)
27 /// .await?;
28 /// ```
29 fn contains(self, lng: f64, lat: f64) -> Select<E>;
30
31 /// 自定义点包含查询 - 查找包含指定点的几何对象
32 ///
33 /// 可以指定数据库字段名和坐标系转换参数
34 ///
35 /// # 参数
36 /// * `lng` - 经度
37 /// * `lat` - 纬度
38 /// * `field` - 数据库中的几何字段名
39 /// * `from_srid` - 源坐标系 SRID
40 /// * `to_srid` - 目标坐标系 SRID
41 ///
42 /// # 示例
43 /// ```
44 /// use my_macros::SpatialQueryExt;
45 ///
46 /// let result = Entity::find()
47 /// .contains_custom(116.26953, 39.74616, "geometry".to_string(), 4326, 4527)
48 /// .all(&db)
49 /// .await?;
50 /// ```
51 fn contains_custom(
52 self,
53 lng: f64,
54 lat: f64,
55 field: String,
56 from_srid: u32,
57 to_srid: u32,
58 ) -> Select<E>;
59
60 /// 多边形包含查询 - 查找完全在指定多边形内的几何对象
61 ///
62 /// 使用默认的 `geom` 字段和 WGS84(4326) 到 CGCS2000(4527) 坐标转换
63 /// 多边形会自动闭合(如果首末点不同)
64 ///
65 /// # 参数
66 /// * `points` - 多边形顶点坐标数组,支持 Vec、数组、切片等类型
67 ///
68 /// # 示例
69 /// ```
70 /// use my_macros::SpatialQueryExt;
71 ///
72 /// let polygon = vec![
73 /// (116.0, 39.0),
74 /// (117.0, 39.0),
75 /// (117.0, 40.0),
76 /// (116.0, 40.0),
77 /// ];
78 ///
79 /// let result = Entity::find()
80 /// .within(polygon)
81 /// .all(&db)
82 /// .await?;
83 /// ```
84 fn within<P>(self, points: P) -> Select<E>
85 where
86 P: AsRef<[(f64, f64)]>;
87
88 /// 自定义多边形包含查询 - 查找完全在指定多边形内的几何对象
89 ///
90 /// 可以指定数据库字段名和坐标系转换参数
91 /// 多边形会自动闭合(如果首末点不同)
92 ///
93 /// # 参数
94 /// * `points` - 多边形顶点坐标数组,支持 Vec、数组、切片等类型
95 /// * `field` - 数据库中的几何字段名
96 /// * `from_srid` - 源坐标系 SRID
97 /// * `to_srid` - 目标坐标系 SRID
98 ///
99 /// # 示例
100 /// ```
101 /// use my_macros::SpatialQueryExt;
102 ///
103 /// let polygon = [(116.0, 39.0), (117.0, 39.0), (117.0, 40.0), (116.0, 40.0)];
104 ///
105 /// let result = Entity::find()
106 /// .within_custom(polygon, "geometry".to_string(), 4326, 4527)
107 /// .all(&db)
108 /// .await?;
109 /// ```
110 fn within_custom<P>(self, points: P, field: String, from_srid: u32, to_srid: u32) -> Select<E>
111 where
112 P: AsRef<[(f64, f64)]>;
113}
114
115impl<E> SpatialQueryExt<E> for Select<E>
116where
117 E: EntityTrait,
118{
119 fn contains(self, lng: f64, lat: f64) -> Select<E> {
120 let expr = Expr::cust(format!(
121 "ST_Contains( geom, ST_Transform( ST_SetSRID(ST_MakePoint({lng}, {lat}), 4326), 4527 ))"
122 ));
123 self.filter(expr)
124 }
125
126 fn contains_custom(
127 self,
128 lng: f64,
129 lat: f64,
130 field: String,
131 from_srid: u32,
132 to_srid: u32,
133 ) -> Select<E> {
134 let expr = Expr::cust(format!(
135 "ST_Contains( {field}, ST_Transform( ST_SetSRID(ST_MakePoint({lng}, {lat}), {from_srid}), {to_srid} ))"
136 ));
137 self.filter(expr)
138 }
139
140 fn within<P>(self, points: P) -> Select<E>
141 where
142 P: AsRef<[(f64, f64)]>,
143 {
144 let wkt = build_polygon_wkt(points.as_ref());
145 let expr = Expr::cust(format!(
146 "ST_Within( geom, ST_Transform( ST_SetSRID(ST_GeomFromText('{wkt}'), 4326), 4527 ))"
147 ));
148 self.filter(expr)
149 }
150
151 fn within_custom<P>(self, points: P, field: String, from_srid: u32, to_srid: u32) -> Select<E>
152 where
153 P: AsRef<[(f64, f64)]>,
154 {
155 let wkt = build_polygon_wkt(points.as_ref());
156 let expr = Expr::cust(format!(
157 "ST_Within( {field}, ST_Transform( ST_SetSRID(ST_GeomFromText('{wkt}'), {from_srid}), {to_srid} ))"
158 ));
159 self.filter(expr)
160 }
161}
162
163/// 构建多边形的 WKT (Well-Known Text) 格式字符串
164///
165/// 将坐标点数组转换为 PostGIS 可识别的多边形 WKT 格式
166/// 自动处理多边形闭合逻辑
167///
168/// # 参数
169/// * `points` - 多边形顶点坐标数组
170///
171/// # 返回值
172/// * 多边形的 WKT 格式字符串,如 "POLYGON((116 39, 117 39, 117 40, 116 40, 116 39))"
173/// * 如果输入为空则返回 "POLYGON EMPTY"
174fn build_polygon_wkt(points: &[(f64, f64)]) -> String {
175 if points.is_empty() {
176 return "POLYGON EMPTY".to_string();
177 }
178
179 let mut coordinates = Vec::new();
180 for (lng, lat) in points {
181 coordinates.push(format!("{lng} {lat}"));
182 }
183
184 if let (Some(first), Some(last)) = (points.first(), points.last()) {
185 if first != last {
186 let (first_lng, first_lat) = first;
187 coordinates.push(format!("{first_lng} {first_lat}"));
188 }
189 }
190
191 format!("POLYGON(({}))", coordinates.join(", "))
192}