Skip to content

Commit ab4801d

Browse files
chat2db-jerrychat2db-jerry
committed
Complete the datasource code
1 parent be87d9e commit ab4801d

29 files changed

Lines changed: 444 additions & 42 deletions

File tree

‎chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/enums/RoleCodeEnum.java‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,23 +13,25 @@ public enum RoleCodeEnum implements BaseEnum<String> {
1313
/**
1414
* DESKTOP
1515
*/
16-
DESKTOP("DESKTOP"),
16+
DESKTOP("DESKTOP", 1L),
1717

1818
/**
19-
* USER
19+
* ADMIN
2020
*/
21-
USER("USER"),
21+
ADMIN("ADMIN", 2L),
2222

2323
/**
24-
* ADMIN
24+
* USER
2525
*/
26-
ADMIN("ADMIN"),
26+
USER("USER", null),
2727

2828
;
2929
final String description;
30+
final Long defaultUserId;
3031

31-
RoleCodeEnum(String description) {
32+
RoleCodeEnum(String description, Long defaultUserId) {
3233
this.description = description;
34+
this.defaultUserId = defaultUserId;
3335
}
3436

3537
@Override

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

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,21 +2,19 @@
22

33
import java.util.List;
44

5-
import jakarta.validation.constraints.NotNull;
6-
75
import ai.chat2db.server.domain.api.model.DataSource;
86
import ai.chat2db.server.domain.api.param.DataSourceCreateParam;
97
import ai.chat2db.server.domain.api.param.DataSourcePageQueryParam;
108
import ai.chat2db.server.domain.api.param.DataSourcePreConnectParam;
119
import ai.chat2db.server.domain.api.param.DataSourceSelector;
1210
import ai.chat2db.server.domain.api.param.DataSourceUpdateParam;
13-
import ai.chat2db.spi.model.Database;
1411
import ai.chat2db.server.tools.base.wrapper.result.ActionResult;
1512
import ai.chat2db.server.tools.base.wrapper.result.DataResult;
1613
import ai.chat2db.server.tools.base.wrapper.result.ListResult;
1714
import ai.chat2db.server.tools.base.wrapper.result.PageResult;
18-
19-
import com.jcraft.jsch.JSchException;
15+
import ai.chat2db.server.tools.common.exception.PermissionDeniedBusinessException;
16+
import ai.chat2db.spi.model.Database;
17+
import jakarta.validation.constraints.NotNull;
2018

2119
/**
2220
* 数据源管理服务
@@ -76,13 +74,24 @@ public interface DataSourceService {
7674
*/
7775
PageResult<DataSource> queryPage(DataSourcePageQueryParam param, DataSourceSelector selector);
7876

77+
/**
78+
* 分页查询数据源列表
79+
* Need to determine permissions
80+
*
81+
* @param param
82+
* @param selector
83+
* @return
84+
* @throws PermissionDeniedBusinessException
85+
*/
86+
PageResult<DataSource> queryPageWithPermission(DataSourcePageQueryParam param, DataSourceSelector selector);
87+
7988
/**
8089
* 通过ID列表查询数据源
8190
*
8291
* @param ids
8392
* @return
8493
*/
85-
ListResult<DataSource> queryByIds(List<Long>ids);
94+
ListResult<DataSource> queryByIds(List<Long> ids);
8695

8796
/**
8897
* 数据源连接测试

‎chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/cache/CacheKey.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@
44

55
public class CacheKey {
66

7+
public static String getLoginUserKey(Long userId) {
8+
return "login_user_" + userId;
9+
}
10+
711
public static String getDataSourceKey(Long dataSourceId) {
812
return "schemas_datasourceId_" + dataSourceId;
913
}

‎chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/cache/MemoryCacheManage.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
package ai.chat2db.server.domain.core.cache;
2+
3+
import java.io.Serializable;
4+
import java.time.Duration;
5+
import java.util.function.Supplier;
6+
7+
import org.apache.commons.lang3.SerializationUtils;
8+
import org.apache.commons.lang3.StringUtils;
9+
import org.ehcache.Cache;
10+
import org.ehcache.config.builders.CacheConfigurationBuilder;
11+
import org.ehcache.config.builders.CacheManagerBuilder;
12+
import org.ehcache.config.builders.ExpiryPolicyBuilder;
13+
import org.ehcache.config.builders.ResourcePoolsBuilder;
14+
import org.ehcache.config.units.MemoryUnit;
15+
import org.springframework.cache.support.NullValue;
16+
17+
/**
18+
* It will only be stored in memory
19+
*
20+
* @author Jiaju Zhuang
21+
*/
22+
public class MemoryCacheManage {
23+
24+
private static final byte[] NULL_BYTES = SerializationUtils.serialize((NullValue)NullValue.INSTANCE);
25+
private static final String CACHE = "memory_cache";
26+
private static final String SYNCHRONIZED_PREFIX = "MemoryCache:";
27+
28+
private static Cache<String, byte[]> cache;
29+
30+
static {
31+
cache = CacheManagerBuilder.newCacheManagerBuilder()
32+
.build(true)
33+
.createCache(CACHE,
34+
CacheConfigurationBuilder.newCacheConfigurationBuilder(String.class, byte[].class,
35+
ResourcePoolsBuilder.newResourcePoolsBuilder()
36+
.offheap(10, MemoryUnit.MB))
37+
.withExpiry(ExpiryPolicyBuilder.timeToLiveExpiration(Duration.ofMinutes(10))));
38+
}
39+
40+
/**
41+
* Retrieve a value from the cache, and if not, query it
42+
* The timeout is fixed at 10 minutes
43+
*
44+
* @param key
45+
* @param queryData
46+
* @param <T>
47+
* @return
48+
*/
49+
public static <T extends Serializable> T computeIfAbsent(String key, Supplier<T> queryData) {
50+
if (key == null) {
51+
return null;
52+
}
53+
T data = get(key);
54+
if (data != null) {
55+
return data;
56+
}
57+
String lockKey = SYNCHRONIZED_PREFIX + key;
58+
synchronized (lockKey.intern()) {
59+
data = get(key);
60+
if (data != null) {
61+
return data;
62+
}
63+
64+
T value = queryData.get();
65+
put(key, value);
66+
return value;
67+
}
68+
}
69+
70+
/**
71+
* Get a data from cache
72+
*
73+
* @param key
74+
* @param <T>
75+
* @return
76+
*/
77+
public static <T> T get(String key) {
78+
if (StringUtils.isBlank(key)) {
79+
return null;
80+
}
81+
byte[] bytes = cache.get(key);
82+
if (bytes == null) {
83+
return null;
84+
}
85+
T data = SerializationUtils.deserialize(bytes);
86+
if (NullValue.INSTANCE.equals(data)) {
87+
return null;
88+
}
89+
return data;
90+
}
91+
92+
/**
93+
* Put a data from cache
94+
* The timeout is fixed at 10 minutes
95+
*
96+
* @param key
97+
* @param value
98+
*/
99+
public static void put(String key, Serializable value) {
100+
if (key == null) {
101+
return;
102+
}
103+
if (value == null) {
104+
cache.put(key, NULL_BYTES);
105+
} else {
106+
cache.put(key, SerializationUtils.serialize(value));
107+
}
108+
}
109+
110+
}

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

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,14 @@
1818
import ai.chat2db.server.domain.api.service.DatabaseService;
1919
import ai.chat2db.server.domain.core.converter.DataSourceConverter;
2020
import ai.chat2db.server.domain.repository.entity.DataSourceDO;
21+
import ai.chat2db.server.domain.repository.mapper.DataSourceCustomMapper;
2122
import ai.chat2db.server.domain.repository.mapper.DataSourceMapper;
2223
import ai.chat2db.server.tools.base.wrapper.result.ActionResult;
2324
import ai.chat2db.server.tools.base.wrapper.result.DataResult;
2425
import ai.chat2db.server.tools.base.wrapper.result.ListResult;
2526
import ai.chat2db.server.tools.base.wrapper.result.PageResult;
27+
import ai.chat2db.server.tools.common.model.LoginUser;
28+
import ai.chat2db.server.tools.common.util.ContextUtils;
2629
import ai.chat2db.spi.config.DriverConfig;
2730
import ai.chat2db.spi.model.DataSourceConnect;
2831
import ai.chat2db.spi.model.Database;
@@ -31,9 +34,10 @@
3134
import ai.chat2db.spi.sql.IDriverManager;
3235
import ai.chat2db.spi.sql.SQLExecutor;
3336
import ai.chat2db.spi.util.JdbcUtils;
34-
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
37+
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
3538
import com.baomidou.mybatisplus.core.metadata.IPage;
3639
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
40+
import jakarta.annotation.Resource;
3741
import lombok.extern.slf4j.Slf4j;
3842
import org.apache.commons.lang3.BooleanUtils;
3943
import org.apache.commons.lang3.StringUtils;
@@ -57,12 +61,15 @@ public class DataSourceServiceImpl implements DataSourceService {
5761

5862
@Autowired
5963
private DatabaseService databaseService;
64+
@Resource
65+
private DataSourceCustomMapper dataSourceCustomMapper;
6066

6167
@Override
6268
public DataResult<Long> create(DataSourceCreateParam param) {
6369
DataSourceDO dataSourceDO = dataSourceConverter.param2do(param);
6470
dataSourceDO.setGmtCreate(LocalDateTime.now());
6571
dataSourceDO.setGmtModified(LocalDateTime.now());
72+
dataSourceDO.setUserId(ContextUtils.getUserId());
6673
dataSourceMapper.insert(dataSourceDO);
6774
preWarmingData(dataSourceDO.getId());
6875
return DataResult.of(dataSourceDO.getId());
@@ -120,9 +127,11 @@ public DataResult<Long> copyById(Long id) {
120127

121128
@Override
122129
public PageResult<DataSource> queryPage(DataSourcePageQueryParam param, DataSourceSelector selector) {
123-
QueryWrapper<DataSourceDO> queryWrapper = new QueryWrapper<>();
130+
LambdaQueryWrapper<DataSourceDO> queryWrapper = new LambdaQueryWrapper<>();
124131
if (StringUtils.isNotBlank(param.getSearchKey())) {
125-
queryWrapper.like("alias", param.getSearchKey());
132+
queryWrapper.and(wrapper -> wrapper.like(DataSourceDO::getAlias, "%" + param.getSearchKey() + "%")
133+
.or()
134+
.like(DataSourceDO::getUrl, "%" + param.getSearchKey() + "%"));
126135
}
127136
Integer start = param.getPageNo();
128137
Integer offset = param.getPageSize();
@@ -132,6 +141,17 @@ public PageResult<DataSource> queryPage(DataSourcePageQueryParam param, DataSour
132141
return PageResult.of(dataSources, iPage.getTotal(), param);
133142
}
134143

144+
@Override
145+
public PageResult<DataSource> queryPageWithPermission(DataSourcePageQueryParam param, DataSourceSelector selector) {
146+
LoginUser loginUser = ContextUtils.getLoginUser();
147+
IPage<DataSourceDO> iPage = dataSourceCustomMapper.selectPageWithPermission(
148+
new Page<>(param.getPageNo(), param.getPageSize()),
149+
BooleanUtils.isTrue(loginUser.getAdmin()), loginUser.getId(), param.getSearchKey());
150+
List<DataSource> dataSources = dataSourceConverter.do2dto(iPage.getRecords());
151+
return PageResult.of(dataSources, iPage.getTotal(), param);
152+
153+
}
154+
135155
@Override
136156
public ListResult<DataSource> queryByIds(List<Long> ids) {
137157
List<DataSourceDO> dataSourceDOS = dataSourceMapper.selectBatchIds(ids);

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package ai.chat2db.server.domain.repository.mapper;
2+
3+
import ai.chat2db.server.domain.repository.entity.DataSourceDO;
4+
import com.baomidou.mybatisplus.core.metadata.IPage;
5+
import org.apache.ibatis.annotations.Param;
6+
7+
/**
8+
* Data Source Custom Mapper
9+
*
10+
* @author Jiaju Zhuang
11+
*/
12+
public interface DataSourceCustomMapper {
13+
14+
IPage<DataSourceDO> selectPageWithPermission(IPage<DataSourceDO> page, @Param("admin") Boolean admin, @Param("userId") Long userId,
15+
@Param("searchKey") String searchKey);
16+
}

‎chat2db-server/chat2db-server-domain/chat2db-server-domain-repository/src/main/resources/ai/chat2db/server/domain/repository/DataSourceCustomMapper.xml‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<!DOCTYPE mapper PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN" "http://mybatis.org/dtd/mybatis-3-mapper.dtd">
3+
<mapper namespace="ai.chat2db.server.domain.repository.mapper.DataSourceCustomMapper">
4+
5+
<select id="selectPageWithPermission" resultType="ai.chat2db.server.domain.repository.entity.DataSourceDO">
6+
select ds.*
7+
from DATA_SOURCE ds
8+
<where>
9+
<if test="admin != true">
10+
<choose>
11+
<when test="userId == 1">
12+
ds.USER_ID = #{userId}
13+
</when>
14+
<otherwise>
15+
(ds.USER_ID = #{userId}
16+
or exists(select 1 from DATA_SOURCE_ACCESS dsa where dsa.ACCESS_OBJECT_TYPE = 'USER' and
17+
dsa.ACCESS_OBJECT_ID = #{userId})
18+
or exists(select 1
19+
from DATA_SOURCE_ACCESS dsa
20+
LEFT JOIN TEAM_USER tu on tu.ID = dsa.ACCESS_OBJECT_ID and dsa.ACCESS_OBJECT_TYPE = 'TEAM'
21+
where tu.USER_ID = #{userId})
22+
)
23+
</otherwise>
24+
</choose>
25+
</if>
26+
<if test="searchKey != '' and searchKey != null ">
27+
and (ds.alias like concat('%',#{searchKey},'%') or ds.url like concat('%',#{searchKey},'%'))
28+
</if>
29+
</where>
30+
31+
</select>
32+
</mapper>

‎chat2db-server/chat2db-server-start/src/main/java/ai/chat2db/server/start/Application.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
package ai.chat2db.server.start;
22

3+
import ai.chat2db.server.tools.common.enums.ModeEnum;
34
import ai.chat2db.server.tools.common.model.ConfigJson;
45
import ai.chat2db.server.tools.common.util.ConfigUtils;
6+
import ai.chat2db.server.tools.common.util.EasyEnumUtils;
57
import com.dtflys.forest.springboot.annotation.ForestScan;
68
import lombok.extern.slf4j.Slf4j;
9+
import org.apache.commons.lang3.ArrayUtils;
710
import org.apache.commons.lang3.StringUtils;
811
import org.mybatis.spring.annotation.MapperScan;
912
import org.springframework.boot.SpringApplication;
@@ -34,12 +37,18 @@ public static void main(String[] args) {
3437
String currentVersion = ConfigUtils.getLocalVersion();
3538
ConfigJson configJson = ConfigUtils.getConfig();
3639
// Represents that the current version has been successfully launched
37-
if (StringUtils.isNotBlank(currentVersion) && StringUtils.equals(currentVersion, configJson.getLatestStartupSuccessVersion())) {
40+
if (StringUtils.isNotBlank(currentVersion) && StringUtils.equals(currentVersion,
41+
configJson.getLatestStartupSuccessVersion())) {
3842
// Flyway doesn't need to start every time to increase startup speed
3943
//args = ArrayUtils.add(args, "--spring.flyway.enabled=false");
4044
log.info("The current version {} has been successfully launched once and will no longer load Flyway.",
4145
currentVersion);
4246
}
47+
ModeEnum mode = EasyEnumUtils.getEnum(ModeEnum.class, System.getProperty("chat2db.mode"));
48+
if (mode == ModeEnum.DESKTOP) {
49+
// In this mode, no user login is required, so only local access is available
50+
args = ArrayUtils.add(args, "--server.address = 0.0.0.0");
51+
}
4352
SpringApplication.run(Application.class, args);
4453
}
4554
}

0 commit comments

Comments
 (0)