SimpleQueryCounter.java
/*
** Module : SimpleQueryCounter.java
** Abstract : Simple query counter implementing the query counter MBean
**
** Copyright (c) 2024, Golden Code Development Corporation.
**
** -#- -I- --Date-- ---------------------------------Description----------------------------------
** 001 ICP 20241114 Created initial version.
*/
/*
** This program is free software: you can redistribute it and/or modify
** it under the terms of the GNU Affero General Public License as
** published by the Free Software Foundation, either version 3 of the
** License, or (at your option) any later version.
**
** This program is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
** GNU Affero General Public License for more details.
**
** You may find a copy of the GNU Affero GPL version 3 at the following
** location: https://www.gnu.org/licenses/agpl-3.0.en.html
**
** Additional terms under GNU Affero GPL version 3 section 7:
**
** Under Section 7 of the GNU Affero GPL version 3, the following additional
** terms apply to the works covered under the License. These additional terms
** are non-permissive additional terms allowed under Section 7 of the GNU
** Affero GPL version 3 and may not be removed by you.
**
** 0. Attribution Requirement.
**
** You must preserve all legal notices or author attributions in the covered
** work or Appropriate Legal Notices displayed by works containing the covered
** work. You may not remove from the covered work any author or developer
** credit already included within the covered work.
**
** 1. No License To Use Trademarks.
**
** This license does not grant any license or rights to use the trademarks
** Golden Code, FWD, any Golden Code or FWD logo, or any other trademarks
** of Golden Code Development Corporation. You are not authorized to use the
** name Golden Code, FWD, or the names of any author or contributor, for
** publicity purposes without written authorization.
**
** 2. No Misrepresentation of Affiliation.
**
** You may not represent yourself as Golden Code Development Corporation or FWD.
**
** You may not represent yourself for publicity purposes as associated with
** Golden Code Development Corporation, FWD, or any author or contributor to
** the covered work, without written authorization.
**
** 3. No Misrepresentation of Source or Origin.
**
** You may not represent the covered work as solely your work. All modified
** versions of the covered work must be marked in a reasonable way to make it
** clear that the modified work is not originating from Golden Code Development
** Corporation or FWD. All modified versions must contain the notices of
** attribution required in this license.
*/
package com.goldencode.p2j.jmx;
import com.goldencode.p2j.util.logging.CentralLogger;
import java.io.*;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.logging.Level;
import java.util.regex.*;
import java.util.stream.Collectors;
public class SimpleQueryCounter
implements SimpleQueryCounterMBean,
QueryCounter
{
/** Logger */
private static final CentralLogger LOG = CentralLogger.get(QueryProfilerJMX.class);
/** Supplier for the initial configuration. */
private Function<String, QueryCountConfiguration> supplier;
/** List of registered databases, set via the bootstrap method. */
private List<String> registeredDatabases = new ArrayList<>();
/** Map that includes the counter configuration for each given database. */
private Map<String, QueryCountConfiguration> configs = new ConcurrentHashMap<>();
/** Map that includes the count of all read queries for each database. */
private final Map<String, Integer> dbReadQueryCountMap = new ConcurrentHashMap<>();
/** Map that includes the count of all write queries for each database. */
private final Map<String, Integer> dbWriteQueryCountMap = new ConcurrentHashMap<>();
/** Nested map that includes the count of read queries for each table in each database. */
private final Map<String, Map<String, Integer>> dbReadTableQueryCountMap = new ConcurrentHashMap<>();
/** Nested map that includes the count of write queries for each table in each database. */
private final Map<String, Map<String, Integer>> dbWriteTableQueryCountMap = new ConcurrentHashMap<>();
/** Holds the total count of read queries across all databases. */
private volatile int readQueriesCount;
/** Holds the total count of write queries across all databases. */
private volatile int writeQueriesCount;
/**
* Getter for the readQueriesCount attribute
*
* @return The total read query count
*/
@Override
public int getReadQueriesCount()
{
return readQueriesCount;
}
/**
* Getter for the writeQueriesCount attribute
*
* @return The total read query count
*/
@Override
public int getWriteQueriesCount()
{
return writeQueriesCount;
}
/**
* Updates the configuration for a specific database.
*
* @param databaseName
* The name of the database.
* @param config
* The configuration object containing enable state and file path.
*/
@Override
public void updateConfiguration(String databaseName, QueryCountConfiguration config)
{
configs.put(databaseName, config);
}
/**
* Updates the query count for a specific database, for each table and increments the total count.
*
* @param databaseName
* The database name for which the query count is updated.
* @param sqlInfo
* The QueryInfo object containing query information
*/
@Override
public void updateCount(String databaseName, QueryInfo sqlInfo)
{
QueryCountConfiguration config = configs.get(databaseName);
if (config != null && config.isEnabled())
{
String sql = sqlInfo.getSql();
if (sql == null)
{
sql = sqlInfo.getReceiver().toString();
}
if (sql == null)
{
return;
}
// Detect tables in the query
List<String> tableNames = extractTableNamesFromSQL(sql);
if (isReadQuery(sql))
{
// Update total read query count per database
dbReadQueryCountMap.merge(databaseName, 1, Integer::sum);
// Update read query count for each table
dbReadTableQueryCountMap.computeIfAbsent(databaseName, db -> new ConcurrentHashMap<>());
for (String tableName : tableNames)
{
dbReadTableQueryCountMap.get(databaseName).merge(tableName, 1, Integer::sum);
}
readQueriesCount ++;
}
else if (isWriteQuery(sql))
{
// Update total write query count per database
dbWriteQueryCountMap.merge(databaseName, 1, Integer::sum);
// Update write query count for each table
dbWriteTableQueryCountMap.computeIfAbsent(databaseName, db -> new ConcurrentHashMap<>());
for (String tableName : tableNames)
{
dbWriteTableQueryCountMap.get(databaseName).merge(tableName, 1, Integer::sum);
}
writeQueriesCount ++;
}
}
}
/**
* Bootstrap method to register all databases with the query counter.
*
* @param databases
* The list of database names to register.
*/
@Override
public void bootstrap(List<String> databases)
{
if (databases == null || databases.isEmpty())
{
LOG.warning("No databases provided for bootstrap.");
return;
}
registeredDatabases = new ArrayList<>(databases);
if (supplier != null)
{
for (String database : databases)
{
QueryCountConfiguration config = supplier.apply(database);
if (config != null)
{
configs.put(database, config);
LOG.info("Configuration for database '" + database + "' has been initialized.");
}
else
{
LOG.info("No configuration found for database '" + database + "' from supplier.");
}
}
}
else
{
LOG.warning("No supplier provided for initializing database configurations.");
}
LOG.info("Registered databases: " + registeredDatabases);
}
/**
* Method to register the supplier for initial configuration.
*
* @param supplier
* The supplier to register.
*/
@Override
public void updateSupplier(Function<String, QueryCountConfiguration> supplier)
{
this.supplier = supplier;
}
/**
* Enables query counting for a specific database.
*
* @param databaseName
* The database name for which counting is enabled.
*/
@Override
public void startCounting(String databaseName)
{
QueryCountConfiguration config = configs.get(databaseName);
if (config == null)
{
config = new QueryCountConfiguration(true, null);
configs.put(databaseName, config);
}
else
{
config.setEnabled(true);
}
}
/**
* Stops query counting for a specific database and generates the report.
*
* @param databaseName
* The database name.
* @param filepath
* The file path to write the report.
*/
@Override
public void stopCounting(String databaseName, String filepath)
{
QueryCountConfiguration config = configs.get(databaseName);
if (config != null)
{
if (isFilePathValid(filepath))
{
config.setFilepath(filepath);
}
generateReportAtFilePath(filepath, Collections.singletonList(databaseName));
config.setEnabled(false);
}
}
/**
* Enables query counting for all databases.
*/
@Override
public void startCountingAll()
{
if (registeredDatabases.isEmpty())
{
LOG.warning("No databases registered for query counting.");
return;
}
registeredDatabases.forEach(
database ->
{
QueryCountConfiguration config = configs.getOrDefault(database,
new QueryCountConfiguration(true,
null));
config.setEnabled(true);
}
);
LOG.log(Level.INFO, "Query counting started for all databases.");
}
/**
* Stops query counting for all databases, updating file paths if valid.
*
* @param filepath
* The file path to write all reports.
*/
@Override
public void stopCountingAll(String filepath)
{
if (registeredDatabases.isEmpty())
{
LOG.warning("No databases registered.");
return;
}
boolean validFilePath = isFilePathValid(filepath);
if (validFilePath)
{
registeredDatabases.forEach(
database ->
{
QueryCountConfiguration config = configs.getOrDefault(database,
new QueryCountConfiguration(true,
null));
config.setFilepath(filepath);
}
);
}
else
{
LOG.severe("Provided path is invalid! Query counting stopped for all databases and " +
"report generated at default file path from directory!");
}
generateReport(registeredDatabases.stream().map(database -> database).collect(Collectors.toList()));
registeredDatabases.forEach(
database ->
{
QueryCountConfiguration config = configs.getOrDefault(database,
new QueryCountConfiguration(false,
null));
config.setEnabled(false);
}
);
LOG.info("Query counting stopped for all databases and report generated at: " + filepath + ".");
}
/**
* Resets query counting for a specific database.
*
* @param databaseName
* The database name.
*/
@Override
public void reset(String databaseName)
{
dbReadQueryCountMap.remove(databaseName);
dbWriteQueryCountMap.remove(databaseName);
dbReadQueryCountMap.remove(databaseName);
dbWriteQueryCountMap.remove(databaseName);
//Compute the new values for the total counters
readQueriesCount = dbReadQueryCountMap.values().stream().mapToInt(Integer::intValue).sum();
writeQueriesCount = dbWriteQueryCountMap.values().stream().mapToInt(Integer::intValue).sum();
LOG.info( "Query counting reset for " + databaseName + ".");
}
/**
* Resets query counting for all databases.
*/
@Override
public void resetAll()
{
dbReadQueryCountMap.clear();
dbWriteQueryCountMap.clear();
dbReadTableQueryCountMap.clear();
dbWriteTableQueryCountMap.clear();
readQueriesCount = 0;
writeQueriesCount = 0;
LOG.log(Level.INFO, "Query counting reset for all databases.");
}
/**
* Generates a report for a list of databases.
*
* @param databases
* The list of database names.
*/
public void generateReport(List<String> databases)
{
Map<String, List<String>> fileToDatabasesMap = new HashMap<>();
// Populate the map with databases grouped by their output file path
for (String databaseName : databases)
{
QueryCountConfiguration config = configs.get(databaseName);
if (config != null && config.isEnabled())
{
String outputFile = config.getFilepath();
if (outputFile != null)
{
fileToDatabasesMap.computeIfAbsent(outputFile, k -> new ArrayList<>()).add(databaseName);
}
}
}
// Generate reports for each unique file path with the associated databases
fileToDatabasesMap.forEach(this::generateReportAtFilePath);
}
/**
* Generates a report for a specific output file, aggregating data from multiple databases.
*
* @param filePath
* The file path to write the report.
* @param databases
* List of databases to include in the report.
*/
private void generateReportAtFilePath(String filePath, List<String> databases)
{
try (FileWriter writer = new FileWriter(filePath))
{
// Write the header for the CSV report
writer.write("Database,Total Read Queries,Total Write Queries," +
"Table,Table Read Queries,Table Write Queries\n");
for (String databaseName : databases)
{
QueryCountConfiguration config = configs.get(databaseName);
if (config == null || !config.isEnabled())
{
LOG.info("Query counting is disabled for " + databaseName + " database.");
continue;
}
// Get total read and write query counts for the database
Integer totalReadQueries = dbReadQueryCountMap.getOrDefault(databaseName, 0);
Integer totalWriteQueries = dbWriteQueryCountMap.getOrDefault(databaseName, 0);
// Write the total read/write query counts for the database
writer.write(databaseName + "," + totalReadQueries + "," + totalWriteQueries + ",,,\n");
// Get table-level read/write query counts
Map<String, Integer> readTableCounts =
dbReadTableQueryCountMap.getOrDefault(databaseName, new ConcurrentHashMap<>());
Map<String, Integer> writeTableCounts =
dbWriteTableQueryCountMap.getOrDefault(databaseName, new ConcurrentHashMap<>());
// Combine all tables from read and write maps
Set<String> allTables = new HashSet<>(readTableCounts.keySet());
allTables.addAll(writeTableCounts.keySet());
// Write table-level read/write query counts
for (String tableName : allTables)
{
Integer tableReadCount = readTableCounts.getOrDefault(tableName, 0);
Integer tableWriteCount = writeTableCounts.getOrDefault(tableName, 0);
writer.write(",,," + tableName + "," + tableReadCount + "," + tableWriteCount + "\n");
}
}
writer.flush();
}
catch (IOException e)
{
LOG.severe("Error writing to file: " + filePath, e);
}
}
/**
* Determines if the SQL is a read query (SELECT).
*
* @param sqlOrReceiver
* The SQL string or receiver string to check.
* @return True if the query is a SELECT, false otherwise.
*/
private boolean isReadQuery(String sqlOrReceiver)
{
String actualSql = extractActualSql(sqlOrReceiver);
return actualSql.trim().toUpperCase().startsWith("SELECT");
}
/**
* Determines if the SQL is a write query (INSERT/DELETE/UPDATE).
*
* @param sqlOrReceiver
* The SQL string or receiver string to check.
* @return True if the query is an INSERT/DELETE/UPDATE, false otherwise.
*/
private boolean isWriteQuery(String sqlOrReceiver)
{
String actualSql = extractActualSql(sqlOrReceiver);
String trimmedSql = actualSql.trim().toUpperCase();
return trimmedSql.startsWith("INSERT") ||
trimmedSql.startsWith("UPDATE") ||
trimmedSql.startsWith("DELETE");
}
/**
* Extracts the actual SQL statement from the string representation of a
* PreparedStatement or SQL string, removing unwanted prefixes like "prepXXXX:".
*
* @param sqlOrReceiver
* The string containing the SQL or PreparedStatement representation.
* @return The cleaned and extracted SQL statement.
*/
private String extractActualSql(String sqlOrReceiver)
{
if (sqlOrReceiver == null || sqlOrReceiver.isEmpty())
{
return "";
}
// Check if the string contains "[wrapping:" which is indicative of the detailed representation
int wrappingIndex = sqlOrReceiver.indexOf("[wrapping:");
if (wrappingIndex != -1)
{
// Extract the part after "[wrapping:" and trim to get the SQL
int startIndex = sqlOrReceiver.indexOf(":", wrappingIndex) + 2;
int endIndex = sqlOrReceiver.lastIndexOf("]");
if (startIndex > 0 && endIndex > startIndex)
{
String rawSql = sqlOrReceiver.substring(startIndex, endIndex).trim();
// Remove any "prepXXXX:" prefix if present
return rawSql.replaceFirst("^prep\\d+:\\s*", "").trim();
}
}
// Fallback for cases where no "[wrapping:" is detected
// Look for "prepXXXX:" and extract the SQL
int prepIndex = sqlOrReceiver.indexOf("prep");
if (prepIndex != -1)
{
int startIndex = sqlOrReceiver.indexOf(":", prepIndex) + 2;
if (startIndex > 0)
{
return sqlOrReceiver.substring(startIndex).replaceFirst("^prep\\d+:\\s*", "").trim();
}
}
// If neither "[wrapping:" nor "prepXXXX:" is found, return the input as is
return sqlOrReceiver;
}
/**
* Checks if a specified match position is inside a brace-enclosed section following a VALUES clause.
* This method searches for a "VALUES(" pattern, followed by a closing parenthesis and an opening brace.
* It then ensures that the match position is within the scope defined by the brace pair.
*
* @param sql
* The SQL string to check.
* @param matchPosition
* The position of the match to check.
* @return {@code true} if the match position is inside braces following a VALUES clause;
* {@code false} otherwise.
*/
private boolean isInsideValuesBraces(String sql, int matchPosition)
{
// Find "VALUES (" and its closing parenthesis
int valuesPosition = sql.toUpperCase().lastIndexOf("VALUES(", matchPosition);
int closingParen = findClosingParen(sql, valuesPosition + "VALUES(".length());
if (valuesPosition == -1)
{
valuesPosition = sql.toUpperCase().lastIndexOf("VALUES (", matchPosition);
closingParen = findClosingParen(sql, valuesPosition + "VALUES (".length());
}
if (valuesPosition == -1)
{
return false;
}
if (closingParen == -1 || closingParen > matchPosition)
{
return false;
}
// Find the first opening "{" after the closing parenthesis of VALUES(...)
int openingBrace = sql.indexOf("{", closingParen);
if (openingBrace == -1 || openingBrace > matchPosition)
{
return false;
}
// Find the corresponding closing brace "}" for the first "{"
int closingBrace = findMatchingBrace(sql, openingBrace);
return closingBrace != -1 && openingBrace < matchPosition && closingBrace > matchPosition;
}
/**
* Finds the position of the closing parenthesis that matches an opening parenthesis
* at a given position, handling nested parentheses.
*
* @param sql
* The SQL string to search within.
* @param openParenPos
* The position of the opening parenthesis.
* @return The position of the matching closing parenthesis,
* {@code -1} if not found.
*/
private int findClosingParen(String sql, int openParenPos)
{
int balance = 1;
for (int i = openParenPos; i < sql.length(); i++)
{
if (sql.charAt(i) == '(') balance++;
if (sql.charAt(i) == ')') balance--;
if (balance == 0)
{
return i;
}
}
return -1;
}
/**
* Finds the position of the closing brace that matches an opening brace
* at a given position, handling nested braces.
*
* @param sql
* The SQL string to search within.
* @param openBracePos
* The position of the opening brace.
* @return The position of the matching closing brace,
* {@code -1} if not found.
*/
private int findMatchingBrace(String sql, int openBracePos)
{
int balance = 1;
for (int i = openBracePos + 1; i < sql.length(); i++)
{
if (sql.charAt(i) == '{') balance++;
if (sql.charAt(i) == '}') balance--;
if (balance == 0)
{
return i;
}
}
return -1;
}
/**
* Extracts unique table names from an SQL query by identifying keywords such as {@code FROM},
* {@code JOIN}, {@code UPDATE}, {@code INSERT INTO}, {@code DELETE FROM}, and subquery patterns.
* Uses regular expressions to locate table names after these clauses, ensuring they are outside
* {@code VALUES} braces, to avoid capturing irrelevant matches.
*
* @param sql
* The SQL query string from which to extract table names.
*
* @return A list of unique table names extracted from the SQL string.
*/
private List<String> extractTableNamesFromSQL(String sql)
{
List<String> tables = new ArrayList<>();
if (sql == null || sql.isEmpty())
{
return tables;
}
String regex = "(?i)\\bFROM\\b\\s+([a-zA-Z_][\\w]*)" +
"|\\bJOIN\\b\\s+([a-zA-Z_][\\w]*)" +
"|\\bUPDATE\\b\\s+([a-zA-Z_][\\w]*)" +
"|\\bINSERT\\s+INTO\\b\\s+([a-zA-Z_][\\w]*)" +
"|\\bDELETE\\s+FROM\\b\\s+([a-zA-Z_][\\w]*)" +
"|\\bWITH\\s+(?:\\w+\\s+AS\\s*\\()([a-zA-Z_][\\w]*)" +
"|\\b\\(SELECT\\s+.*?\\s+FROM\\s+([a-zA-Z_][\\w]*)";
Pattern pattern = Pattern.compile(regex);
Matcher matcher = pattern.matcher(sql);
while (matcher.find())
{
for (int i = 1; i <= matcher.groupCount(); i++)
{
String table = matcher.group(i);
if (table != null &&
!isInsideValuesBraces(sql, matcher.start(i)) &&
!tables.contains(table))
{
tables.add(table);
}
}
}
return tables;
}
/**
* Checks if the provided file path is valid.
*
* @param filepath
* The file path to check.
* @return true if the path is valid; false otherwise.
*/
private boolean isFilePathValid(String filepath)
{
File file = new File(filepath);
if (file.exists())
{
return file.canWrite();
}
File parent = file.getParentFile();
return (parent != null ? parent.canWrite() : new File(".").canWrite());
}
}