diff --git a/mateclaw-server/src/main/java/vip/mate/datasource/service/DatasourceService.java b/mateclaw-server/src/main/java/vip/mate/datasource/service/DatasourceService.java index 1c72747d..01295d34 100644 --- a/mateclaw-server/src/main/java/vip/mate/datasource/service/DatasourceService.java +++ b/mateclaw-server/src/main/java/vip/mate/datasource/service/DatasourceService.java @@ -109,6 +109,7 @@ public class DatasourceService { // 更新测试结果 entity.setLastTestTime(LocalDateTime.now()); entity.setLastTestOk(ok); + encryptPassword(entity); datasourceMapper.updateById(entity); return ok; } diff --git a/mateclaw-server/src/main/java/vip/mate/tool/builtin/DatasourceTool.java b/mateclaw-server/src/main/java/vip/mate/tool/builtin/DatasourceTool.java index 9c464a7b..bc4bd650 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/builtin/DatasourceTool.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/builtin/DatasourceTool.java @@ -110,9 +110,11 @@ public class DatasourceTool { "SELECT TABLE_NAME, TABLE_COMMENT, TABLE_ROWS FROM information_schema.TABLES WHERE TABLE_SCHEMA = '%s' ORDER BY TABLE_NAME", sanitizeIdentifier(entity.getDatabaseName())); case "postgresql" -> String.format( - "SELECT tablename AS table_name, obj_description(c.oid) AS table_comment " + - "FROM pg_tables t LEFT JOIN pg_class c ON c.relname = t.tablename " + - "WHERE t.schemaname = '%s' ORDER BY tablename", + "SELECT t.table_name, obj_description(c.oid) AS table_comment " + + "FROM information_schema.tables t " + + "LEFT JOIN pg_namespace n ON n.nspname = t.table_schema " + + "LEFT JOIN pg_class c ON c.relnamespace = n.oid AND c.relname = t.table_name " + + "WHERE t.table_schema = '%s' ORDER BY t.table_name", sanitizeIdentifier(entity.getSchemaName() != null ? entity.getSchemaName() : "public")); case "clickhouse" -> "SHOW TABLES"; default -> throw new IllegalArgumentException("不支持的数据库类型: " + dbType); diff --git a/mateclaw-server/src/test/java/vip/mate/datasource/service/DatasourceServicePasswordPersistenceTest.java b/mateclaw-server/src/test/java/vip/mate/datasource/service/DatasourceServicePasswordPersistenceTest.java new file mode 100644 index 00000000..6ca99239 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/datasource/service/DatasourceServicePasswordPersistenceTest.java @@ -0,0 +1,48 @@ +package vip.mate.datasource.service; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.test.util.ReflectionTestUtils; +import vip.mate.datasource.model.DatasourceEntity; +import vip.mate.datasource.repository.DatasourceMapper; + +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class DatasourceServicePasswordPersistenceTest { + + @Test + @DisplayName("connection test never persists a decrypted datasource password") + void testConnectionReencryptsPasswordBeforeUpdate() { + // Given a datasource whose password is encrypted at creation time. + String plaintext = "reader-password"; + DatasourceMapper mapper = mock(DatasourceMapper.class); + DatasourceConnectionManager connectionManager = mock(DatasourceConnectionManager.class); + DatasourceService service = new DatasourceService(mapper, connectionManager); + ReflectionTestUtils.setField(service, "encryptKey", "test-datasource-key"); + + DatasourceEntity datasource = new DatasourceEntity(); + datasource.setId(1L); + datasource.setName("DUCHA"); + datasource.setPassword(plaintext); + service.create(datasource); + assertNotEquals(plaintext, datasource.getPassword()); + when(mapper.selectById(1L)).thenReturn(datasource); + doAnswer(invocation -> { + assertTrue(plaintext.equals(datasource.getPassword())); + return true; + }).when(connectionManager).testConnection(datasource); + + // When MateClaw tests the JDBC connection. + assertTrue(service.testConnection(1L)); + + // Then the entity sent back to MyBatis must be encrypted again. + assertNotEquals(plaintext, datasource.getPassword()); + assertTrue(datasource.getPassword().matches("[0-9a-f]+")); + verify(mapper).updateById(datasource); + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/tool/builtin/DatasourceToolPostgresqlViewDiscoveryTest.java b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DatasourceToolPostgresqlViewDiscoveryTest.java new file mode 100644 index 00000000..ba2b48a0 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DatasourceToolPostgresqlViewDiscoveryTest.java @@ -0,0 +1,64 @@ +package vip.mate.tool.builtin; + +import com.fasterxml.jackson.databind.json.JsonMapper; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import vip.mate.datasource.model.DatasourceEntity; +import vip.mate.datasource.service.DatasourceConnectionManager; +import vip.mate.datasource.service.DatasourceService; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class DatasourceToolPostgresqlViewDiscoveryTest { + + @Test + @DisplayName("list_tables discovers PostgreSQL views as queryable relations") + void listTablesIncludesPostgresqlViews() throws Exception { + // Given a PostgreSQL datasource whose allowed schema exposes a view. + DatasourceEntity datasource = new DatasourceEntity(); + datasource.setId(1L); + datasource.setDbType("postgresql"); + datasource.setSchemaName("serving"); + + DatasourceService service = mock(DatasourceService.class); + DatasourceConnectionManager connectionManager = mock(DatasourceConnectionManager.class); + Connection connection = mock(Connection.class); + Statement statement = mock(Statement.class); + ResultSet resultSet = mock(ResultSet.class); + ResultSetMetaData metadata = mock(ResultSetMetaData.class); + when(service.getDecrypted(1L)).thenReturn(datasource); + when(connectionManager.getConnection(datasource)).thenReturn(connection); + when(connection.createStatement()).thenReturn(statement); + when(statement.executeQuery(anyString())).thenReturn(resultSet); + when(resultSet.getMetaData()).thenReturn(metadata); + when(metadata.getColumnCount()).thenReturn(2); + when(metadata.getColumnLabel(1)).thenReturn("table_name"); + when(metadata.getColumnLabel(2)).thenReturn("table_comment"); + when(resultSet.next()).thenReturn(true, false); + when(resultSet.getString(1)).thenReturn("v_plan_enriched"); + when(resultSet.getString(2)).thenReturn("read-only serving view"); + + DatasourceTool tool = new DatasourceTool(service, connectionManager, JsonMapper.builder().build()); + + // When the agent asks MateClaw to discover available relations. + String output = tool.query_datasource("list_tables", 1L, null); + + // Then the metadata query must include views and return the discovered view. + ArgumentCaptor sql = ArgumentCaptor.forClass(String.class); + verify(statement).executeQuery(sql.capture()); + assertTrue(sql.getValue().contains("information_schema.tables")); + assertFalse(sql.getValue().contains("FROM pg_tables")); + assertTrue(output.contains("v_plan_enriched")); + } +}