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}