Add logic for verifying connection permissions

chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/DataSourceAccessBusinessService.java

@@ -0,0 +1,19 @@
+package ai.chat2db.server.domain.api.service;
+
+import ai.chat2db.server.tools.base.wrapper.result.ActionResult;
+import jakarta.validation.constraints.NotNull;
+
+/**
+ * Data Source Access
+ *
+ * @author Jiaju Zhuang
+ */
+public interface DataSourceAccessBusinessService {
+    /**
+     * delete
+     *
+     * @param dataSourceId
+     * @return
+     */
+    ActionResult checkPermission(@NotNull Long dataSourceId);
+}

chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DataSourceAccessBusinessServiceImpl.java

@@ -0,0 +1,65 @@
+package ai.chat2db.server.domain.core.impl;
+
+import ai.chat2db.server.domain.api.enums.AccessObjectTypeEnum;
+import ai.chat2db.server.domain.api.enums.RoleCodeEnum;
+import ai.chat2db.server.domain.api.param.datasource.access.DataSourceAccessPageQueryParam;
+import ai.chat2db.server.domain.api.service.DataSourceAccessBusinessService;
+import ai.chat2db.server.domain.api.service.DataSourceAccessService;
+import ai.chat2db.server.domain.repository.mapper.DataSourceAccessCustomMapper;
+import ai.chat2db.server.tools.base.wrapper.result.ActionResult;
+import ai.chat2db.server.tools.common.exception.PermissionDeniedBusinessException;
+import ai.chat2db.server.tools.common.model.LoginUser;
+import ai.chat2db.server.tools.common.util.ContextUtils;
+import jakarta.annotation.Resource;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.stereotype.Service;
+
+/**
+ * Data Source Access
+ *
+ * @author Jiaju Zhuang
+ */
+@Slf4j
+@Service
+public class DataSourceAccessBusinessServiceImpl implements DataSourceAccessBusinessService {
+
+    @Resource
+    private DataSourceAccessService dataSourceAccessService;
+    @Resource
+    private DataSourceAccessCustomMapper dataSourceAccessCustomMapper;
+
+    @Override
+    public ActionResult checkPermission(Long dataSourceId) {
+        LoginUser loginUser = ContextUtils.getLoginUser();
+        // Representative is desktop mode
+        if (RoleCodeEnum.DESKTOP.getDefaultUserId().equals(loginUser.getId())) {
+            if (RoleCodeEnum.DESKTOP.getDefaultUserId().equals(dataSourceId)) {
+                return ActionResult.isSuccess();
+            } else {
+                throw new PermissionDeniedBusinessException();
+            }
+        }
+
+        // Administrators can edit anything
+        if (loginUser.getAdmin()) {
+            return ActionResult.isSuccess();
+        }
+
+        // Verify if user have permission
+        DataSourceAccessPageQueryParam dataSourceAccessPageQueryParam = new DataSourceAccessPageQueryParam();
+        dataSourceAccessPageQueryParam.setDataSourceId(dataSourceId);
+        dataSourceAccessPageQueryParam.setAccessObjectType(AccessObjectTypeEnum.USER.getCode());
+        dataSourceAccessPageQueryParam.setAccessObjectId(loginUser.getId());
+        dataSourceAccessPageQueryParam.queryOne();
+        if (dataSourceAccessService.pageQuery(dataSourceAccessPageQueryParam, null).hasData()) {
+            return ActionResult.isSuccess();
+        }
+
+        // Verify if the team has permission
+        if (dataSourceAccessCustomMapper.checkTeamPermission(dataSourceId, loginUser.getId()) != null) {
+            return ActionResult.isSuccess();
+
+        }
+        throw new PermissionDeniedBusinessException();
+    }
+}

chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DataSourceServiceImpl.java

@@ -117,7 +117,7 @@ public class DataSourceServiceImpl implements DataSourceService {
     @Override
     public ActionResult updateWithPermission(DataSourceUpdateParam param) {
         DataSource dataSource = queryExistent(param.getId()).getData();
-        PermissionUtils.checkPermission(dataSource.getUserId());
+        PermissionUtils.checkOperationPermission(dataSource.getUserId());
 
         DataSourceDO dataSourceDO = dataSourceConverter.param2do(param);
         dataSourceDO.setGmtModified(LocalDateTime.now());
@@ -129,7 +129,7 @@ public class DataSourceServiceImpl implements DataSourceService {
     public ActionResult deleteWithPermission(Long id) {
 
         DataSource dataSource = queryExistent(id).getData();
-        PermissionUtils.checkPermission(dataSource.getUserId());
+        PermissionUtils.checkOperationPermission(dataSource.getUserId());
 
         dataSourceMapper.deleteById(id);
         return ActionResult.isSuccess();
@@ -153,7 +153,7 @@ public class DataSourceServiceImpl implements DataSourceService {
     @Override
     public DataResult<Long> copyByIdWithPermission(Long id) {
         DataSource dataSource = queryExistent(id).getData();
-        PermissionUtils.checkPermission(dataSource.getUserId());
+        PermissionUtils.checkOperationPermission(dataSource.getUserId());
 
         DataSourceDO dataSourceDO = dataSourceMapper.selectById(id);
         dataSourceDO.setId(null);

chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/util/PermissionUtils.java

@@ -17,7 +17,7 @@ public class PermissionUtils {
      *
      * @param createUserId The creator of the current content
      */
-    public static void checkPermission(Long createUserId) {
+    public static void checkOperationPermission(Long createUserId) {
         LoginUser loginUser = ContextUtils.getLoginUser();
         // Representative is desktop mode
         if (RoleCodeEnum.DESKTOP.getDefaultUserId().equals(loginUser.getId())) {
@@ -36,4 +36,24 @@ public class PermissionUtils {
             throw new PermissionDeniedBusinessException();
         }
     }
+
+    /**
+     * 校验是否有查询权限
+     *
+     * @param createUserId
+     * @return
+     */
+    public static boolean checkBaseQueryPermission(Long createUserId) {
+        LoginUser loginUser = ContextUtils.getLoginUser();
+        // Representative is desktop mode
+        if (RoleCodeEnum.DESKTOP.getDefaultUserId().equals(loginUser.getId())) {
+            if (RoleCodeEnum.DESKTOP.getDefaultUserId().equals(createUserId)) {
+                return true;
+            } else {
+                throw new PermissionDeniedBusinessException();
+            }
+        }
+        // Administrators can edit anything
+        return loginUser.getAdmin();
+    }
 }

chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/java/ai/chat2db/server/domain/repository/mapper/DataSourceAccessCustomMapper.java

@@ -17,4 +17,6 @@ public interface DataSourceAccessCustomMapper extends Mapper<DataSourceAccessDO>
         @Param("accessObjectId") Long accessObjectId,
         @Param("userOrTeamSearchKey") String userOrTeamSearchKey,
         @Param("dataSourceSearchKey") String dataSourceSearchKey);
+
+    DataSourceAccessDO checkTeamPermission( @Param("dataSourceId") Long dataSourceId, @Param("userId") Long userId);
 }

chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/mapper/DataSourceAccessCustomMapper.xml

@@ -35,4 +35,17 @@
         </where>
     </select>
 
+    <select id="checkTeamPermission" resultType="ai.chat2db.server.domain.repository.entity.DataSourceAccessDO">
+        select dsa.*
+        from DATA_SOURCE_ACCESS dsa
+        left join TEAM t on t.id =dsa.ACCESS_OBJECT_ID and dsa.ACCESS_OBJECT_TYPE = 'TEAM' and t.status='VALID'
+        left join TEAM_USER tu on tu.TEAM_ID =t.ID
+        <where>
+            dsa.DATA_SOURCE_ID = #{dataSourceId}
+            and tu.USER_ID = #{userId}
+            limit 1
+        </where>
+    </select>
+
+
 </mapper>

chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/aspect/ConnectionInfoHandler.java

@@ -2,6 +2,7 @@
 package ai.chat2db.server.web.api.aspect;
 
 import ai.chat2db.server.domain.api.model.DataSource;
+import ai.chat2db.server.domain.api.service.DataSourceAccessBusinessService;
 import ai.chat2db.server.domain.api.service.DataSourceService;
 import ai.chat2db.server.tools.base.wrapper.result.DataResult;
 import ai.chat2db.server.tools.common.exception.ParamBusinessException;
@@ -11,6 +12,7 @@ import ai.chat2db.server.web.api.controller.data.source.request.DataSourceConsol
 import ai.chat2db.spi.config.DriverConfig;
 import ai.chat2db.spi.sql.Chat2DBContext;
 import ai.chat2db.spi.sql.ConnectInfo;
+import jakarta.annotation.Resource;
 import lombok.extern.slf4j.Slf4j;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
@@ -29,6 +31,8 @@ public class ConnectionInfoHandler {
 
     @Autowired
     private DataSourceService dataSourceService;
+    @Resource
+    private DataSourceAccessBusinessService dataSourceAccessBusinessService;
 
     @Around("within(@ai.chat2db.server.web.api.aspect.ConnectionInfoAspect *)")
     public Object connectionInfoHandler(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {
@@ -37,16 +41,16 @@ public class ConnectionInfoHandler {
             if (params != null && params.length > 0) {
                 for (int i = 0; i < params.length; i++) {
                     Object param = params[i];
-                    if(param instanceof DataSourceBaseRequest){
+                    if (param instanceof DataSourceBaseRequest) {
                         Long dataSourceId = ((DataSourceBaseRequest)param).getDataSourceId();
                         String schemaName = ((DataSourceBaseRequest)param).getSchemaName();
                         String database = ((DataSourceBaseRequest)param).getDatabaseName();
-                        Chat2DBContext.putContext(toInfo(dataSourceId, database, null,schemaName));
+                        Chat2DBContext.putContext(toInfo(dataSourceId, database, null, schemaName));
-                    }else if (param instanceof DataSourceConsoleRequestInfo) {
+                    } else if (param instanceof DataSourceConsoleRequestInfo) {
                         Long dataSourceId = ((DataSourceConsoleRequestInfo)param).getDataSourceId();
                         Long consoleId = ((DataSourceConsoleRequestInfo)param).getConsoleId();
                         String database = ((DataSourceConsoleRequestInfo)param).getDatabaseName();
-                        Chat2DBContext.putContext(toInfo(dataSourceId, database, consoleId,null));
+                        Chat2DBContext.putContext(toInfo(dataSourceId, database, consoleId, null));
                     } else if (param instanceof DataSourceBaseRequestInfo) {
                         Long dataSourceId = ((DataSourceBaseRequestInfo)param).getDataSourceId();
                         String database = ((DataSourceBaseRequestInfo)param).getDatabaseName();
@@ -60,12 +64,16 @@ public class ConnectionInfoHandler {
         }
     }
 
-    public ConnectInfo toInfo(Long dataSourceId, String database, Long consoleId,String schemaName) {
+    public ConnectInfo toInfo(Long dataSourceId, String database, Long consoleId, String schemaName) {
         DataResult<DataSource> result = dataSourceService.queryById(dataSourceId);
         DataSource dataSource = result.getData();
         if (!result.success() || dataSource == null) {
             throw new ParamBusinessException("dataSourceId");
         }
+
+        // Verify permissions
+        dataSourceAccessBusinessService.checkPermission(dataSourceId);
+
         ConnectInfo connectInfo = new ConnectInfo();
         connectInfo.setAlias(dataSource.getAlias());
         connectInfo.setUser(dataSource.getUserName());
@@ -93,6 +101,7 @@ public class ConnectionInfoHandler {
     }
 
     public ConnectInfo toInfo(Long dataSourceId, String database) {
-        return toInfo(dataSourceId, database, null,null);
+        return toInfo(dataSourceId, database, null, null);
     }
+
 }