diff --git a/agentscope-extensions/agentscope-extensions-postgresql/src/main/java/io/agentscope/extensions/postgresql/store/PostgresBaseStore.java b/agentscope-extensions/agentscope-extensions-postgresql/src/main/java/io/agentscope/extensions/postgresql/store/PostgresBaseStore.java index fca7ea358..d530d850a 100644 --- a/agentscope-extensions/agentscope-extensions-postgresql/src/main/java/io/agentscope/extensions/postgresql/store/PostgresBaseStore.java +++ b/agentscope-extensions/agentscope-extensions-postgresql/src/main/java/io/agentscope/extensions/postgresql/store/PostgresBaseStore.java @@ -84,7 +84,7 @@ public class PostgresBaseStore implements BaseStore { /** Default table name used when the builder is not customised. */ public static final String DEFAULT_TABLE_NAME = "agentscope_store"; - private static final Pattern VALID_TABLE_NAME = Pattern.compile("[A-Za-z_][A-Za-z0-9_]*"); + private static final Pattern VALID_IDENTIFIER = Pattern.compile("[A-Za-z_][A-Za-z0-9_]*"); private static final String CREATE_TABLE_SQL = """ @@ -139,7 +139,9 @@ ON CONFLICT (namespace_path, item_key) DO UPDATE SET private final DataSource dataSource; private final ObjectMapper objectMapper; + private final String schemaName; private final String tableName; + private final String qualifiedTableName; private final String selectSql; private final String upsertSql; @@ -151,13 +153,15 @@ ON CONFLICT (namespace_path, item_key) DO UPDATE SET private PostgresBaseStore(Builder b) { this.dataSource = b.dataSource; this.objectMapper = b.objectMapper != null ? b.objectMapper : new ObjectMapper(); + this.schemaName = b.schemaName; this.tableName = b.tableName; - this.selectSql = String.format(SELECT_SQL, tableName); - this.upsertSql = String.format(UPSERT_SQL, tableName); - this.insertSql = String.format(INSERT_SQL, tableName); - this.casUpdateSql = String.format(CAS_UPDATE_SQL, tableName); - this.deleteSql = String.format(DELETE_SQL, tableName); - this.searchSql = String.format(SEARCH_SQL, tableName); + this.qualifiedTableName = getQualifiedTableName(schemaName, tableName); + this.selectSql = String.format(SELECT_SQL, qualifiedTableName); + this.upsertSql = String.format(UPSERT_SQL, qualifiedTableName); + this.insertSql = String.format(INSERT_SQL, qualifiedTableName); + this.casUpdateSql = String.format(CAS_UPDATE_SQL, qualifiedTableName); + this.deleteSql = String.format(DELETE_SQL, qualifiedTableName); + this.searchSql = String.format(SEARCH_SQL, qualifiedTableName); if (b.initializeSchema) { initializeSchema(); } @@ -169,14 +173,25 @@ public static Builder builder(DataSource dataSource) { } private void initializeSchema() { - String ddl = String.format(CREATE_TABLE_SQL, tableName); + String ddl = String.format(CREATE_TABLE_SQL, qualifiedTableName); try (Connection c = dataSource.getConnection(); Statement st = c.createStatement()) { + if (schemaName != null) { + st.executeUpdate("CREATE SCHEMA IF NOT EXISTS \"" + schemaName + "\""); + } st.executeUpdate(ddl); } catch (SQLException e) { throw new IllegalStateException( - "Failed to initialize PostgresBaseStore schema for table " + tableName, e); + "Failed to initialize PostgresBaseStore schema for table " + qualifiedTableName, + e); + } + } + + private static String getQualifiedTableName(String schemaName, String tableName) { + if (schemaName == null) { + return tableName; } + return "\"" + schemaName + "\".\"" + tableName + "\""; } @Override @@ -367,6 +382,7 @@ public static final class Builder { private final DataSource dataSource; private ObjectMapper objectMapper; + private String schemaName; private String tableName = DEFAULT_TABLE_NAME; private boolean initializeSchema; @@ -381,18 +397,25 @@ public Builder objectMapper(ObjectMapper objectMapper) { } /** - * Overrides the table name. Must match the regex {@code [A-Za-z_][A-Za-z0-9_]*}; this is - * the only place a table identifier ever flows into the SQL string verbatim, so the - * validation here is the SQL-injection guard. + * Sets the PostgreSQL schema name to use. If not set, SQL uses the unqualified table name + * for backwards compatibility. + * + *
The provided schema name must be non-null, non-blank, and match {@code [A-Za-z_][A-Za-z0-9_]*}. + * The schema is created when {@link #initializeSchema(boolean)} is enabled. + */ + public Builder schemaName(String schemaName) { + this.schemaName = validateIdentifier(schemaName, "schemaName"); + return this; + } + + /** + * Overrides the table name. Identifiers must match the regex {@code [A-Za-z_][A-Za-z0-9_]*}. + * + *
Schema and table identifiers are embedded into SQL (identifier quoting is used when a + * schema is configured), so validation here is the SQL-injection guard. */ public Builder tableName(String tableName) { - if (tableName == null - || tableName.isBlank() - || !VALID_TABLE_NAME.matcher(tableName).matches()) { - throw new IllegalArgumentException( - "tableName must match [A-Za-z_][A-Za-z0-9_]*, got: " + tableName); - } - this.tableName = tableName; + this.tableName = validateIdentifier(tableName, "tableName"); return this; } @@ -409,8 +432,18 @@ public Builder initializeSchema(boolean initializeSchema) { /** Builds the {@link PostgresBaseStore}. */ public PostgresBaseStore build() { PostgresBaseStore store = new PostgresBaseStore(this); - LOG.debug("PostgresBaseStore built: table={}", store.tableName); + LOG.debug("PostgresBaseStore built: table={}", store.qualifiedTableName); return store; } + + private static String validateIdentifier(String identifier, String identifierName) { + if (identifier == null + || identifier.isBlank() + || !VALID_IDENTIFIER.matcher(identifier).matches()) { + throw new IllegalArgumentException( + identifierName + " must match [A-Za-z_][A-Za-z0-9_]*, got: " + identifier); + } + return identifier; + } } } diff --git a/agentscope-extensions/agentscope-extensions-postgresql/src/test/java/io/agentscope/extensions/postgresql/store/PostgresBaseStoreTest.java b/agentscope-extensions/agentscope-extensions-postgresql/src/test/java/io/agentscope/extensions/postgresql/store/PostgresBaseStoreTest.java index 471b08468..be461b9cd 100644 --- a/agentscope-extensions/agentscope-extensions-postgresql/src/test/java/io/agentscope/extensions/postgresql/store/PostgresBaseStoreTest.java +++ b/agentscope-extensions/agentscope-extensions-postgresql/src/test/java/io/agentscope/extensions/postgresql/store/PostgresBaseStoreTest.java @@ -95,6 +95,19 @@ void builderWithCustomTableName() { assertNotNull(store); } + @Test + void builderWithCustomSchemaNameUsesQualifiedTableName() throws SQLException { + PostgresBaseStore store = + PostgresBaseStore.builder(dataSource).schemaName("custom_schema").build(); + + store.get(List.of("namespace"), "key"); + + verify(connection) + .prepareStatement( + "SELECT value_json, version FROM \"custom_schema\".\"agentscope_store\"" + + " WHERE namespace_path = ? AND item_key = ?"); + } + @Test void builderRejectsInvalidTableName() { assertThrows( @@ -108,6 +121,19 @@ void builderRejectsInvalidTableName() { () -> PostgresBaseStore.builder(dataSource).tableName(null).build()); } + @Test + void builderRejectsInvalidSchemaName() { + assertThrows( + IllegalArgumentException.class, + () -> PostgresBaseStore.builder(dataSource).schemaName("schema;drop").build()); + assertThrows( + IllegalArgumentException.class, + () -> PostgresBaseStore.builder(dataSource).schemaName("").build()); + assertThrows( + IllegalArgumentException.class, + () -> PostgresBaseStore.builder(dataSource).schemaName(null).build()); + } + @Test void initializeSchemaDefaultDoesNotCreateTable() throws SQLException { PostgresBaseStore store = PostgresBaseStore.builder(dataSource).build(); @@ -123,6 +149,23 @@ void initializeSchemaTrueCreatesTable() throws SQLException { verify(connection).createStatement(); } + @Test + void initializeSchemaTrueCreatesCustomSchemaAndTable() throws SQLException { + PostgresBaseStore store = + PostgresBaseStore.builder(dataSource) + .schemaName("custom_schema") + .initializeSchema(true) + .build(); + + assertNotNull(store); + verify(statement).executeUpdate("CREATE SCHEMA IF NOT EXISTS \"custom_schema\""); + verify(statement) + .executeUpdate( + org.mockito.ArgumentMatchers.contains( + "CREATE TABLE IF NOT EXISTS" + + " \"custom_schema\".\"agentscope_store\"")); + } + @Test void initializeSchemaFailureThrows() throws SQLException { when(connection.createStatement()).thenThrow(new SQLException("boom"));