【java】mybatis拦截器实现单表对于geometry类型字段的新增查询更新

核心思想就是将geometry类型的数据通过pg函数st_asgeojson st_asewkt st_astext 将sql修改。
直接上代码：

1. 放在geometry类型字段上的注解：

   import java.lang.annotation.*;
   
   /**
    * @description: geometry字段的注解 表示该字段是geometry类型的 该注解返回的是java.lang.String
    * @author: QinLei
    * @email: [email protected]
    * @date: 2023年09月12日  10:00
    */
   @Documented
   @Inherited
   @Target({ElementType.FIELD})
   @Retention(RetentionPolicy.RUNTIME)
   public @interface GeometryField {
   
   	/**
   	 * 查询时返回的格式
   	 *
   	 * @return
   	 * @author: QinLei
   	 * @date: 2023年09月12日  10:02
   	 */
   	GeometryType returnFormat() default GeometryType.FEATURE;
   
   	/**
   	 * 对应数据库列名
   	 *
   	 * @return
   	 * @author: QinLei
   	 * @date: 2023年09月12日  10:55
   	 */
   	String colName();
   }
   
   

2. GeometryType 类

   	/**
    * @description: geometry返回的格式
    * 进行修改和新增时，仅支持geojson格式的数据（强制要求设置crs）
    * @author: QinLei
    * @email: [email protected]
    * @date: 2023年09月12日  10:04
    */
   public enum GeometryType {
   
   	/**
   	 * geojson
   	 *
   	 * @author: QinLei
   	 * @date: 2023年09月12日  10:18
   	 */
   	GEOJSON(),
   	/**
   	 * feature
   	 *
   	 * @author: QinLei
   	 * @date: 2023年09月12日  10:19
   	 */
   	FEATURE,
   	/**
   	 * wkt
   	 *
   	 * @author: QinLei
   	 * @date: 2023年09月12日  10:19
   	 */
   	WKT,
   	/**
   	 * text
   	 *
   	 * @author: QinLei
   	 * @date: 2023年09月12日  10:19
   	 */
   	TEXT;
   
   	/**
   	 * 获取当前类型返回的sql
   	 * 修改查询sql
   	 * @param col 数据库表对应列名
   	 * @param tableName 表名
   	 * @param pk 主键字段
   	 * @return
   	 * @author: QinLei
   	 * @date: 2023年09月12日  11:09
   	 */
   	public String getSelectSql(String col, String pk, String tableName) {
   		String sql = new String();
   		switch (this) {
   			case GEOJSON:
   				sql = "jsonb_build_object(\n" +
   					"    'type',       'Feature',\n" +
   					"    'id',         '" + tableName + ".' || " + pk + ",\n" +
   					"    'geometry',   ST_AsGeoJSON(" + col + ")::jsonb,\n" +
   					"    'properties', to_jsonb( " + tableName + ".* ) - '" + col + "'\n" +
   					"    )";
   				break;
   			case FEATURE:
   				sql = "st_asgeojson(" + col + ")";
   				break;
   			case WKT:
   				sql = "st_asewkt(" + col + ")";
   				break;
   			case TEXT:
   				sql = "st_astext(" + col + ")";
   				break;
   			default:
   				sql = col;
   		}
   		return sql;
   	}
   	/**
   	 * 更新和插入sql修改，前端传递geojson相对来说用turf好实现点，所以只支持了geojson
   	 * @return
   	 * @author: QinLei
   	 * @date: 2023年09月15日  10:58
   	 */
   	public String getOtherSql() {
   		return "st_geomfromgeojson(?)";
   	}
   }
   

3. 一些工具类

   import com.alibaba.druid.sql.ast.SQLStatement;
   import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
   import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlSchemaStatVisitor;
   import com.alibaba.druid.sql.parser.SQLStatementParser;
   import com.alibaba.druid.stat.TableStat;
   
   import java.util.ArrayList;
   import java.util.List;
   import java.util.Map;
   
   /**
    * @description: 表名获取工具
    * @author: QinLei
    * @email: [email protected]
    * @date: 2023年09月12日  14:24
    */
   public class TableNameUtil {
   
   	/**
   	 * 根据sql获取表名
   	 *
   	 * @param sql
   	 * @return
   	 * @author: QinLei
   	 * @date: 2023年09月12日  14:24
   	 */
   	public static List<String> getAllTableNameBySQL(String sql) {
   		SQLStatementParser parser = new MySqlStatementParser(sql);
   		// 使用Parser解析生成AST，这里SQLStatement就是AST
   		SQLStatement sqlStatement = parser.parseStatement();
   		MySqlSchemaStatVisitor visitor = new MySqlSchemaStatVisitor();
   		sqlStatement.accept(visitor);
   		Map<TableStat.Name, TableStat> tables = visitor.getTables();
   		List<String> allTableName = new ArrayList<>();
   		for (TableStat.Name t : tables.keySet()) {
   			allTableName.add(t.getName());
   		}
   		return allTableName;
   	}
   }
   
   

   /**
   	 * 替换第n个字符
   	 *
   	 * @param str         要操作的字符串
   	 * @param replacement 要替换成的字符
   	 * @param source      源字符
   	 * @param n           第几个源字符
   	 * @return
   	 * @author: QinLei
   	 * @date: 2023年09月12日  15:35
   	 */
   	public static String replaceNthOccurrence(String str, int n, String source, String replacement) {
   		StringBuilder sb = new StringBuilder(str);
   		int count = 0;
   		int index = 0;
   		while ((index = sb.indexOf(source, index)) != -1) {
   			if (count == n) {
   				String prefix = sb.substring(0, index);
   				String suffix = sb.substring(index + 1);
   				sb.setLength(0);
   				sb.append(prefix).append(replacement).append(suffix);
   				break;
   			}
   			count++;
   			index++;
   		}
   		return sb.toString();
   	}
   

   还有一些简单的判空工具，这里就不列了。

4. 拦截器

   import com.baomidou.mybatisplus.annotation.TableId;
   import org.apache.ibatis.binding.MapperMethod;
   import org.apache.ibatis.executor.statement.StatementHandler;
   import org.apache.ibatis.mapping.*;
   import org.apache.ibatis.plugin.*;
   import org.apache.ibatis.reflection.MetaObject;
   import org.apache.ibatis.reflection.SystemMetaObject;
   import org.apache.ibatis.session.RowBounds;
   import annotation.GeometryField;
   import annotation.GeometryType;
   import utils.TableNameUtil;
   import utils.CollectionUtil;
   import utils.Func;
   import utils.StringUtil;
   import org.springframework.core.annotation.AnnotationUtils;
   import org.springframework.stereotype.Component;
   
   import java.lang.reflect.Field;
   import java.lang.reflect.Proxy;
   import java.sql.Connection;
   import java.util.HashMap;
   import java.util.List;
   import java.util.Map;
   import java.util.Properties;
   
   /**
    * @description: 当前仅支持单表
    * @author: QinLei
    * @email: [email protected]
    * @date: 2023年09月12日  9:51
    */
   @Component
   @Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
   public class GeometryInterceptor implements Interceptor {
   	@Override
   	public Object intercept(Invocation invocation) throws Throwable {
   		StatementHandler statementHandler = (StatementHandler) realTarget(invocation.getTarget());
   		MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
   		MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
   		BoundSql boundSql = (BoundSql) metaObject.getValue("delegate.boundSql");
   		String sql = boundSql.getSql();
   		List<String> allTableNameBySQL = TableNameUtil.getAllTableNameBySQL(sql);
   		// 只支持单表 多表直接跳过
   		if (allTableNameBySQL.size() == 1) {
   			if (SqlCommandType.SELECT.equals(mappedStatement.getSqlCommandType())) {
   				this.selectInvoke(mappedStatement, sql, allTableNameBySQL, metaObject);
   			} else if (SqlCommandType.INSERT.equals(mappedStatement.getSqlCommandType())
   				|| SqlCommandType.UPDATE.equals(mappedStatement.getSqlCommandType())) {
   				this.otherInvoke(boundSql, sql, metaObject, mappedStatement.getSqlCommandType());
   			}
   		}
   		return invocation.proceed();
   	}
   
   
   	/**
   	 * 查询处理
   	 *
   	 * @param mappedStatement
   	 * @param sql
   	 * @param allTableNameBySQL
   	 * @param metaObject
   	 * @author: QinLei
   	 * @date: 2023年09月12日  14:59
   	 */
   	private void selectInvoke(MappedStatement mappedStatement, String sql, List<String> allTableNameBySQL, MetaObject metaObject) throws Exception {
   		List<ResultMap> resultMaps = mappedStatement.getResultMaps();
   		ResultMap map = resultMaps.get(0);
   		Class<?> type = map.getType();
   		Map<String, Object> res = this.selectGeometryMap(type);
   		Map<String, GeometryType> geometryMap = (Map<String, GeometryType>) res.get("geometryMap");
   		StringBuilder sb = new StringBuilder();
   		sb.append("SELECT ");
   		String reSql = sql.replace("SELECT", "");
   		int from = reSql.indexOf("FROM");
   		String subSql = reSql.substring(0, from);
   		String[] colArr = subSql.split(",");
   		for (String col : colArr) {
   			String colName = col.trim();
   			if (geometryMap.containsKey(colName)) {
   				GeometryType geometryType = geometryMap.get(colName);
   				sb.append(geometryType.getSelectSql(colName, String.valueOf(res.get("pk")), allTableNameBySQL.get(0)))
   					.append(" AS ").append(colName).append(",");
   			} else {
   				sb.append(col).append(",");
   			}
   		}
   		String substring = sb.substring(0, sb.length() - 1);
   		String lastSql = reSql.substring(from);
   		metaObject.setValue("delegate.boundSql.sql", substring + " " + lastSql);
   		metaObject.setValue("delegate.rowBounds.offset", RowBounds.NO_ROW_OFFSET);
   		metaObject.setValue("delegate.rowBounds.limit", RowBounds.NO_ROW_LIMIT);
   	}
   
   	/**
   	 * 增删改处理
   	 *
   	 * @param boundSql
   	 * @param sql
   	 * @param metaObject
   	 * @param sqlCommandType
   	 * @author: QinLei
   	 * @date: 2023年09月12日  15:01
   	 */
   	private void otherInvoke(BoundSql boundSql, String sql, MetaObject metaObject, SqlCommandType sqlCommandType) {
   		Object paramObj = null;
   		Object parameterObject = boundSql.getParameterObject();
   		if (sqlCommandType.equals(SqlCommandType.INSERT)) {
   			paramObj = parameterObject;
   		} else {
   			if (parameterObject instanceof MapperMethod.ParamMap) {
   				MapperMethod.ParamMap<?> p = (MapperMethod.ParamMap<?>) parameterObject;
   				if (p.containsKey("et")) {
   					paramObj = p.get("et");
   				} else {
   					paramObj = p.get("param1");
   				}
   			} else {
   				paramObj = parameterObject;
   			}
   		}
   		Map<String, GeometryType> geometryMap = this.insertUpdateGeometryMap(paramObj);
   		List<ParameterMapping> parameterMappings = boundSql.getParameterMappings();
   		int j = 0;
   		for (int i = 0; i < parameterMappings.size(); i++) {
   			ParameterMapping parameterMapping = parameterMappings.get(i);
   			String property = parameterMapping.getProperty();
   			String pv = property;
   			if (property.startsWith("et")) {
   				pv = property.replace("et.", "");
   			}
   			if (geometryMap.containsKey(pv)) {
   				GeometryType geometryType = geometryMap.get(pv);
   				sql = StringUtil.replaceNthOccurrence(sql, j, "?", geometryType.getOtherSql());
   			} else {
   				j++;
   			}
   		}
   		metaObject.setValue("delegate.boundSql.sql", sql);
   	}
   
   	private Map<String, Object> selectGeometryMap(Class<?> type) throws Exception {
   		Map<String, Object> res = new HashMap<>();
   		Map<String, GeometryType> geometryMap = new HashMap<>();
   		Field[] fields = type.getDeclaredFields();
   		String pk = null;
   		for (int i = 0; i < fields.length; i++) {
   			TableId tableId = fields[i].getAnnotation(TableId.class);
   			if (tableId != null) {
   				pk = tableId.value();
   			}
   			GeometryField gf = fields[i].getAnnotation(GeometryField.class);
   			if (!"serialVersionUID".equalsIgnoreCase(fields[i].getName()) && !Func.isNull(gf)) {
   				geometryMap.put(gf.colName(), gf.returnFormat());
   			}
   		}
   		if (StringUtil.isBlank(pk) && CollectionUtil.isNotEmpty(geometryMap)) {
   			throw new Exception("该实体无主键！");
   		}
   		res.put("pk", pk);
   		res.put("geometryMap", geometryMap);
   		return res;
   	}
   
   	private Map<String, GeometryType> insertUpdateGeometryMap(Object paramObj) {
   		Map<String, GeometryType> geometryMap = new HashMap<>();
   		Field[] fields = paramObj.getClass().getDeclaredFields();
   		for (Field declaredField : fields) {
   			declaredField.setAccessible(true);
   			try {
   				Object o = declaredField.get(paramObj);
   				GeometryField geometryField = AnnotationUtils.findAnnotation(declaredField, GeometryField.class);
   				if (o != null && !"serialVersionUID".equalsIgnoreCase(declaredField.getName()) && !Func.isNull(geometryField)) {
   					geometryMap.put(geometryField.colName(), geometryField.returnFormat());
   					geometryMap.put(declaredField.getName(), geometryField.returnFormat());
   				}
   			} catch (IllegalAccessException e) {
   				e.printStackTrace();
   			}
   		}
   		return geometryMap;
   	}
   
   
   	private Object realTarget(Object target) {
   		if (Proxy.isProxyClass(target.getClass())) {
   			MetaObject metaObject = SystemMetaObject.forObject(target);
   			return realTarget(metaObject.getValue("h.target"));
   		}
   		return target;
   	}
   
   
   	@Override
   	public Object plugin(Object target) {
   		return Plugin.wrap(target, this);
   	}
   
   	@Override
   	public void setProperties(Properties properties) {
   
   	}
   
   }