Stamp JDBC backend tables with their tree name and refresh optimizer statistics after import

opendj-server-legacy/src/main/java/org/opends/server/backends/jdbc/JDBCStorage.java

@@ -39,6 +39,7 @@
 import java.security.NoSuchAlgorithmException;
 import java.sql.*;
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 
 import static org.opends.server.backends.pluggable.spi.StorageUtils.addErrorMessage;
 import static org.opends.server.util.StaticUtils.stackTraceToSingleLineString;
@@ -88,6 +89,14 @@ int execute(PreparedStatement statement) throws SQLException {
 		return statement.executeUpdate();
 	}
 
+	// unlike execute(), tolerates statements that return a result set ("analyze table" on mysql)
+	void executeAny(PreparedStatement statement) throws SQLException {
+		if (logger.isTraceEnabled()) {
+			logger.trace(LocalizableMessage.raw("jdbc: %s",statement));
+		}
+		statement.execute();
+	}
+
 	Connection getConnection() throws Exception {
 		return CachedConnection.getConnection(config.getDBDirectory());
 	}
@@ -134,6 +143,182 @@ String getTableName(TreeName treeName) {
 		return tree2table.get(treeName);
 	}
 
+	private static final String[] NO_ARGS=new String[0];
+
+	static String driverNameOf(Connection con) {
+		return ((CachedConnection) con).parent.getClass().getName();
+	}
+
+	// Splices a value into a single-quoted SQL literal for the comment DDL, which takes no bind
+	// parameters: doubles every quote, and every backslash on dialects where backslash is an
+	// escape character inside literals. The scan of the escaped result is defence in depth: it
+	// re-verifies that no quote (or live backslash) is left unpaired and able to terminate the
+	// literal, so a regression in the escaping throws instead of reaching the database.
+	private static String sqlLiteral(String value, boolean backslashIsEscape) {
+		final String escaped=(backslashIsEscape?value.replace("\\","\\\\"):value).replace("'","''");
+		for (int i=0;i<escaped.length();i++) {
+			final char c=escaped.charAt(i);
+			if (c=='\'' || (backslashIsEscape && c=='\\')) {
+				if (i+1>=escaped.length() || escaped.charAt(i+1)!=c) {
+					throw new IllegalArgumentException("unpaired "+c+" in SQL literal: "+escaped);
+				}
+				i++;
+			}
+		}
+		return "'"+escaped+"'";
+	}
+
+	// Table names are opaque SHA-224 hashes, so on the database side there is no way to tell
+	// which tree a table holds. Stamp each table with its tree name (visible in "\dt+" and the
+	// information schema) so database-level troubleshooting does not require recomputing hashes.
+	// Runs on a dedicated pooled connection, never on the transaction that opened the tree:
+	// comment statements are DDL (an implicit commit on mysql and oracle), and a failing
+	// sp_addextendedproperty rolls the whole transaction back on sql server - either would
+	// corrupt work pending on the caller's connection (e.g. the trusted flag written by
+	// DefaultIndex.afterOpen()). Returns true when the comment was stamped; the comment is a
+	// diagnostic aid, so a failed attempt only returns false and must not fail the backend.
+	boolean commentTable(TreeName treeName) {
+		final String tableName=getTableName(treeName);
+		try (final Connection con=getConnection()) {
+			final String driverName=driverNameOf(con);
+			final boolean postgres=driverName.contains("postgres");
+			final boolean mysql=driverName.contains("mysql");
+			final boolean oracle=driverName.contains("oracle");
+			final boolean microsoft=driverName.contains("microsoft");
+			if (!postgres && !mysql && !oracle && !microsoft) { // no comment syntax and readback known for other engines: leave the table unstamped
+				return false;
+			}
+			final String treeComment=treeName.toString();
+			// comment statements are DDL (metadata lock on mysql, ddl lock on oracle) and openTree()
+			// runs on every backend open: only stamp when the stored comment is absent or stale
+			if (treeComment.equals(readStoredComment(con, tableName))) {
+				return false;
+			}
+			final String sql;
+			final String[] args;
+			if (mysql) { // ALTER TABLE takes no binds; backslash is an escape character in mysql literals
+				sql="alter table "+tableName+" comment "+sqlLiteral(treeComment,true);
+				args=NO_ARGS;
+			}else if (microsoft) { // no COMMENT ON in t-sql: MS_Description extended property (procedure arguments take binds)
+				sql="declare @s sysname = schema_name()"
+					+" if exists (select 1 from sys.extended_properties where class=1 and major_id=object_id(?) and minor_id=0 and name='MS_Description')"
+					+" exec sys.sp_updateextendedproperty N'MS_Description', ?, N'SCHEMA', @s, N'TABLE', ?"
+					+" else"
+					+" exec sys.sp_addextendedproperty N'MS_Description', ?, N'SCHEMA', @s, N'TABLE', ?";
+				args=new String[]{tableName, treeComment, tableName, treeComment, tableName};
+			}else if (postgres) { // no binds in ddl; the E'' form keeps backslash an escape character regardless of standard_conforming_strings
+				sql="comment on table "+tableName+" is E"+sqlLiteral(treeComment,true);
+				args=NO_ARGS;
+			}else { // oracle: no binds in ddl; backslash is never an escape character in oracle literals
+				sql="comment on table "+tableName+" is "+sqlLiteral(treeComment,false);
+				args=NO_ARGS;
+			}
+			try (final PreparedStatement statement=con.prepareStatement(sql)) {
+				for (int i=0;i<args.length;i++) {
+					statement.setString(i+1,args[i]);
+				}
+				executeAny(statement);
+				con.commit();
+			}
+			return true;
+		}catch (Exception e) { // CachedConnection.close() has rolled back whatever the failed attempt left pending
+			logger.debug(LocalizableMessage.raw("jdbc: unable to comment table %s with tree name %s: %s",
+				tableName, treeName, stackTraceToSingleLineString(e)));
+			return false;
+		}
+	}
+
+	// Returns the comment currently stored on the table, or null when there is none.
+	// Only called for the four dialects commentTable() recognizes.
+	String readStoredComment(Connection con, String tableName) throws SQLException {
+		final String driverName=driverNameOf(con);
+		final String sql;
+		final String arg;
+		if (driverName.contains("postgres")) {
+			sql="select obj_description(to_regclass(?), 'pg_class')";
+			arg=tableName;
+		}else if (driverName.contains("mysql")) {
+			sql="select table_comment from information_schema.tables where table_schema=database() and table_name=?";
+			arg=tableName;
+		}else if (driverName.contains("oracle")) {
+			sql="select comments from user_tab_comments where table_name=?";
+			arg=tableName.toUpperCase();
+		}else if (driverName.contains("microsoft")) {
+			sql="select cast(value as nvarchar(4000)) from sys.extended_properties where class=1 and major_id=object_id(?) and minor_id=0 and name='MS_Description'";
+			arg=tableName;
+		}else {
+			throw new SQLException("no table comment readback for "+driverName);
+		}
+		try (final PreparedStatement statement=con.prepareStatement(sql)) {
+			statement.setString(1,arg);
+			try (final ResultSet rs=executeResultSet(statement)) {
+				return rs.next() ? rs.getString(1) : null;
+			}
+		}
+	}
+
+	// A bulk load leaves the optimizer statistics of freshly created tables stale (a table that
+	// was never analyzed can make the planner badly misestimate the "where k>? order by k" cursor
+	// batches - see OpenIdentityPlatform/OpenDJ#859), so refresh them once the data is in place.
+	// Only the trees the import actually wrote are refreshed: rebuild-index imports a few index
+	// trees, and gathering statistics of the whole backend on its behalf is a full scan per
+	// table on oracle. Statistics upkeep is best-effort: a failure must not fail the import that
+	// produced the data, so failures are only logged - the return value makes them observable to tests.
+	boolean updateTableStatistics(Connection con, Collection<TreeName> trees) {
+		final String driverName=driverNameOf(con);
+		final boolean postgres=driverName.contains("postgres");
+		final boolean mysql=driverName.contains("mysql");
+		final boolean oracle=driverName.contains("oracle");
+		final boolean microsoft=driverName.contains("microsoft");
+		if (!postgres && !mysql && !oracle && !microsoft) { // no portable statistics refresh for other engines
+			return true;
+		}
+		boolean allRefreshed=true;
+		for (final TreeName treeName : trees) {
+			final String tableName=getTableName(treeName);
+			final String sql;
+			final String[] args;
+			if (postgres) {
+				sql="analyze "+tableName;
+				args=NO_ARGS;
+			}else if (mysql) {
+				sql="analyze table "+tableName;
+				args=NO_ARGS;
+			}else if (oracle) {
+				sql="begin dbms_stats.gather_table_stats(user, ?); end;";
+				args=new String[]{tableName.toUpperCase()};
+			}else {
+				sql="update statistics "+tableName;
+				args=NO_ARGS;
+			}
+			try (final PreparedStatement statement=con.prepareStatement(sql)) {
+				for (int i=0;i<args.length;i++) {
+					statement.setString(i+1,args[i]);
+				}
+				if (mysql) { // mysql reports analyze problems as a result row, not an SQLException
+					try (final ResultSet rs=executeResultSet(statement)) {
+						while (rs.next()) {
+							if ("error".equalsIgnoreCase(rs.getString("Msg_type"))) {
+								throw new SQLException(rs.getString("Msg_text"));
+							}
+						}
+					}
+				}else {
+					executeAny(statement);
+				}
+				con.commit();
+			}catch (SQLException e) {
+				try {
+					con.rollback();
+				} catch (SQLException e2) {}
+				allRefreshed=false;
+				logger.warn(LocalizableMessage.raw("jdbc: unable to refresh statistics of table %s (tree %s): %s",
+					tableName, treeName, stackTraceToSingleLineString(e)));
+			}
+		}
+		return allRefreshed;
+	}
+
 	@Override
 	public void removeStorageFiles() throws StorageRuntimeException {
 		final boolean isOpen=getStorageStatus().isWorking();
@@ -163,6 +348,10 @@ public void removeStorageFiles() throws StorageRuntimeException {
 			} catch (Exception e) {
 				throw new StorageRuntimeException(e);
 			}
+			// all tables are gone: forget the mappings so listTrees() consumers skip the dropped trees
+			for (final TreeName treeName : trees) {
+				tree2table.invalidate(treeName);
+			}
 		}
 		if (!isOpen) {
 			close();
@@ -335,6 +524,7 @@ public void openTree(TreeName treeName, boolean createOnDemand) {
 					}
 				}
 				// mssql: k is varbinary(max), which cannot be an index key column - cursor batches stay unindexed there
+				commentTable(treeName);
 			}
 		}
 
@@ -369,6 +559,8 @@ public void deleteTree(TreeName treeName) {
 					throw new StorageRuntimeException(e);
 				}
 			}
+			// forget the mapping so listTrees() consumers (updateTableStatistics) skip the dropped table
+			tree2table.invalidate(treeName);
 		}
 
 		@Override
@@ -670,9 +862,11 @@ private final class ImporterImpl implements Importer {
 		final Connection con;
 		final ReadableTransactionImpl txr;
 		final WriteableTransactionTransactionImpl txw;
+		// the trees this import wrote: close() refreshes the statistics of these and only these
+		final Set<TreeName> writtenTrees = ConcurrentHashMap.newKeySet();
 
 		final Boolean isOpen;
-		
+
 		public ImporterImpl() {
 			isOpen=getStorageStatus().isWorking();
 			if (!isOpen) {
@@ -694,24 +888,31 @@ public ImporterImpl() {
 		@Override
 		public void close() {
 			try {
-				con.commit();
-				con.close();
+				try {
+					con.commit();
+					updateTableStatistics(con, writtenTrees);
+				} finally { // the pooled connection must be returned even when the commit or a statistics statement throws
+					con.close();
+				}
 			} catch (SQLException e) {
 				throw new StorageRuntimeException(e);
-			}
-			if (!isOpen) {
-				JDBCStorage.this.close();
+			} finally {
+				if (!isOpen) {
+					JDBCStorage.this.close();
+				}
 			}
 		}
-		
+
 		@Override
 		public void clearTree(TreeName name) {
 			txw.clearTree(name);
+			writtenTrees.add(name);
 		}
-		
+
 		@Override
 		public void put(TreeName treeName, ByteSequence key, ByteSequence value) {
 			txw.put(treeName, key, value);
+			writtenTrees.add(treeName);
 		}
 
 		@Override