Skip to content

Commit

Permalink
feat: support client-side hints for tags and priority (#3005)
Browse files Browse the repository at this point in the history 
Supports including a hint in the SQL string to set a statement tag or an RPC priority for a single SQL statement in the Connection API. This makes it easier to use these features from frameworks and tools that only support SQL statements.

Replaces #2978
  • Loading branch information
@olavloite
olavloite committed Apr 19, 2024
1 parent 5419491 commit 48828df
Show file tree
Hide file tree
Showing 11 changed files with 1,303 additions and 124 deletions.
5 changes: 5 additions & 0 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/Statement.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,11 @@ public String getSql() {
return sql;
}

/** Returns a copy of this statement with the SQL string replaced by the given SQL string. */
public Statement withReplacedSql(String sql) {
return new Statement(sql, this.parameters, this.queryOptions);
}

/** Returns the {@link QueryOptions} that will be used with this {@link Statement}. */
public QueryOptions getQueryOptions() {
return queryOptions;
Expand Down
141 changes: 79 additions & 62 deletions ...ud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractStatementParser.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,13 @@

package com.google.cloud.spanner.connection;

import static com.google.cloud.spanner.connection.SimpleParser.isValidIdentifierChar;
import static com.google.cloud.spanner.connection.StatementHintParser.convertHintsToOptions;

import com.google.api.core.InternalApi;
import com.google.cloud.spanner.Dialect;
import com.google.cloud.spanner.ErrorCode;
import com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption;
import com.google.cloud.spanner.SpannerException;
import com.google.cloud.spanner.SpannerExceptionFactory;
import com.google.cloud.spanner.Statement;
Expand Down Expand Up @@ -169,6 +173,7 @@ public static class ParsedStatement {
private final Statement statement;
private final String sqlWithoutComments;
private final boolean returningClause;
private final ReadQueryUpdateTransactionOption[] optionsFromHints;

private static ParsedStatement clientSideStatement(
ClientSideStatementImpl clientSideStatement,
Expand All @@ -182,15 +187,27 @@ private static ParsedStatement ddl(Statement statement, String sqlWithoutComment
}

private static ParsedStatement query(
Statement statement, String sqlWithoutComments, QueryOptions defaultQueryOptions) {
Statement statement,
String sqlWithoutComments,
QueryOptions defaultQueryOptions,
ReadQueryUpdateTransactionOption[] optionsFromHints) {
return new ParsedStatement(
StatementType.QUERY, null, statement, sqlWithoutComments, defaultQueryOptions, false);
StatementType.QUERY,
null,
statement,
sqlWithoutComments,
defaultQueryOptions,
false,
optionsFromHints);
}

private static ParsedStatement update(
Statement statement, String sqlWithoutComments, boolean returningClause) {
Statement statement,
String sqlWithoutComments,
boolean returningClause,
ReadQueryUpdateTransactionOption[] optionsFromHints) {
return new ParsedStatement(
StatementType.UPDATE, statement, sqlWithoutComments, returningClause);
StatementType.UPDATE, statement, sqlWithoutComments, returningClause, optionsFromHints);
}

private static ParsedStatement unknown(Statement statement, String sqlWithoutComments) {
Expand All @@ -208,18 +225,20 @@ private ParsedStatement(
this.statement = statement;
this.sqlWithoutComments = Preconditions.checkNotNull(sqlWithoutComments);
this.returningClause = false;
this.optionsFromHints = EMPTY_OPTIONS;
}

private ParsedStatement(
StatementType type,
Statement statement,
String sqlWithoutComments,
boolean returningClause) {
this(type, null, statement, sqlWithoutComments, null, returningClause);
boolean returningClause,
ReadQueryUpdateTransactionOption[] optionsFromHints) {
this(type, null, statement, sqlWithoutComments, null, returningClause, optionsFromHints);
}

private ParsedStatement(StatementType type, Statement statement, String sqlWithoutComments) {
this(type, null, statement, sqlWithoutComments, null, false);
this(type, null, statement, sqlWithoutComments, null, false, EMPTY_OPTIONS);
}

private ParsedStatement(
Expand All @@ -228,33 +247,37 @@ private ParsedStatement(
Statement statement,
String sqlWithoutComments,
QueryOptions defaultQueryOptions,
boolean returningClause) {
boolean returningClause,
ReadQueryUpdateTransactionOption[] optionsFromHints) {
Preconditions.checkNotNull(type);
this.type = type;
this.clientSideStatement = clientSideStatement;
this.statement = statement == null ? null : mergeQueryOptions(statement, defaultQueryOptions);
this.sqlWithoutComments = Preconditions.checkNotNull(sqlWithoutComments);
this.returningClause = returningClause;
this.optionsFromHints = optionsFromHints;
}

private ParsedStatement copy(Statement statement, QueryOptions defaultQueryOptions) {
return new ParsedStatement(
this.type,
this.clientSideStatement,
statement,
statement.withReplacedSql(this.statement.getSql()),
this.sqlWithoutComments,
defaultQueryOptions,
this.returningClause);
this.returningClause,
this.optionsFromHints);
}

private ParsedStatement forCache() {
return new ParsedStatement(
this.type,
this.clientSideStatement,
null,
Statement.of(this.statement.getSql()),
this.sqlWithoutComments,
null,
this.returningClause);
this.returningClause,
this.optionsFromHints);
}

@Override
Expand Down Expand Up @@ -287,6 +310,11 @@ public boolean hasReturningClause() {
return this.returningClause;
}

@InternalApi
public ReadQueryUpdateTransactionOption[] getOptionsFromHints() {
return this.optionsFromHints;
}

/**
* @return true if the statement is a query that will return a {@link
* com.google.cloud.spanner.ResultSet}.
Expand Down Expand Up @@ -480,14 +508,23 @@ ParsedStatement parse(Statement statement, QueryOptions defaultQueryOptions) {
}

private ParsedStatement internalParse(Statement statement, QueryOptions defaultQueryOptions) {
StatementHintParser statementHintParser =
new StatementHintParser(getDialect(), statement.getSql());
ReadQueryUpdateTransactionOption[] optionsFromHints = EMPTY_OPTIONS;
if (statementHintParser.hasStatementHints()
&& !statementHintParser.getClientSideStatementHints().isEmpty()) {
statement =
statement.toBuilder().replace(statementHintParser.getSqlWithoutClientSideHints()).build();
optionsFromHints = convertHintsToOptions(statementHintParser.getClientSideStatementHints());
}
String sql = removeCommentsAndTrim(statement.getSql());
ClientSideStatementImpl client = parseClientSideStatement(sql);
if (client != null) {
return ParsedStatement.clientSideStatement(client, statement, sql);
} else if (isQuery(sql)) {
return ParsedStatement.query(statement, sql, defaultQueryOptions);
return ParsedStatement.query(statement, sql, defaultQueryOptions, optionsFromHints);
} else if (isUpdateStatement(sql)) {
return ParsedStatement.update(statement, sql, checkReturningClause(sql));
return ParsedStatement.update(statement, sql, checkReturningClause(sql), optionsFromHints);
} else if (isDdlStatement(sql)) {
return ParsedStatement.ddl(statement, sql);
}
Expand Down Expand Up @@ -621,6 +658,10 @@ public String removeCommentsAndTrim(String sql) {
/** Removes any statement hints at the beginning of the statement. */
abstract String removeStatementHint(String sql);

@VisibleForTesting
static final ReadQueryUpdateTransactionOption[] EMPTY_OPTIONS =
new ReadQueryUpdateTransactionOption[0];

/** Parameter information with positional parameters translated to named parameters. */
@InternalApi
public static class ParametersInfo {
Expand Down Expand Up @@ -697,9 +738,10 @@ public boolean checkReturningClause(String sql) {
return checkReturningClauseInternal(sql);
}

abstract Dialect getDialect();

/**
* <<<<<<< HEAD Returns true if this dialect supports nested comments. ======= <<<<<<< HEAD
* Returns true if this dialect supports nested comments. >>>>>>> main
* Returns true if this dialect supports nested comments.
*
* <ul>
* <li>This method should return false for dialects that consider this to be a valid comment:
Expand Down Expand Up @@ -757,18 +799,6 @@ public boolean checkReturningClause(String sql) {
/** Returns the query parameter prefix that should be used for this dialect. */
abstract String getQueryParameterPrefix();

/**
* Returns true for characters that can be used as the first character in unquoted identifiers.
*/
boolean isValidIdentifierFirstChar(char c) {
return Character.isLetter(c) || c == UNDERSCORE;
}

/** Returns true for characters that can be used in unquoted identifiers. */
boolean isValidIdentifierChar(char c) {
return isValidIdentifierFirstChar(c) || Character.isDigit(c) || c == DOLLAR;
}

/** Reads a dollar-quoted string literal from position index in the given sql string. */
String parseDollarQuotedString(String sql, int index) {
// Look ahead to the next dollar sign (if any). Everything in between is the quote tag.
Expand Down Expand Up @@ -812,9 +842,9 @@ int skip(String sql, int currentIndex, @Nullable StringBuilder result) {
} else if (currentChar == HYPHEN
&& sql.length() > (currentIndex + 1)
&& sql.charAt(currentIndex + 1) == HYPHEN) {
return skipSingleLineComment(sql, currentIndex, result);
return skipSingleLineComment(sql, /* prefixLength = */ 2, currentIndex, result);
} else if (currentChar == DASH && supportsHashSingleLineComments()) {
return skipSingleLineComment(sql, currentIndex, result);
return skipSingleLineComment(sql, /* prefixLength = */ 1, currentIndex, result);
} else if (currentChar == SLASH
&& sql.length() > (currentIndex + 1)
&& sql.charAt(currentIndex + 1) == ASTERISK) {
Expand All @@ -826,44 +856,31 @@ int skip(String sql, int currentIndex, @Nullable StringBuilder result) {
}

/** Skips a single-line comment from startIndex and adds it to result if result is not null. */
static int skipSingleLineComment(String sql, int startIndex, @Nullable StringBuilder result) {
int endIndex = sql.indexOf('\n', startIndex + 2);
if (endIndex == -1) {
endIndex = sql.length();
} else {
// Include the newline character.
endIndex++;
int skipSingleLineComment(
String sql, int prefixLength, int startIndex, @Nullable StringBuilder result) {
return skipSingleLineComment(getDialect(), sql, prefixLength, startIndex, result);
}

static int skipSingleLineComment(
Dialect dialect,
String sql,
int prefixLength,
int startIndex,
@Nullable StringBuilder result) {
SimpleParser simpleParser = new SimpleParser(dialect, sql, startIndex, false);
if (simpleParser.skipSingleLineComment(prefixLength)) {
appendIfNotNull(result, sql.substring(startIndex, simpleParser.getPos()));
}
appendIfNotNull(result, sql.substring(startIndex, endIndex));
return endIndex;
return simpleParser.getPos();
}

/** Skips a multi-line comment from startIndex and adds it to result if result is not null. */
int skipMultiLineComment(String sql, int startIndex, @Nullable StringBuilder result) {
// Current position is start + '/*'.length().
int pos = startIndex + 2;
// PostgreSQL allows comments to be nested. That is, the following is allowed:
// '/* test /* inner comment */ still a comment */'
int level = 1;
while (pos < sql.length()) {
if (supportsNestedComments()
&& sql.charAt(pos) == SLASH
&& sql.length() > (pos + 1)
&& sql.charAt(pos + 1) == ASTERISK) {
level++;
}
if (sql.charAt(pos) == ASTERISK && sql.length() > (pos + 1) && sql.charAt(pos + 1) == SLASH) {
level--;
if (level == 0) {
pos += 2;
appendIfNotNull(result, sql.substring(startIndex, pos));
return pos;
}
}
pos++;
SimpleParser simpleParser = new SimpleParser(getDialect(), sql, startIndex, false);
if (simpleParser.skipMultiLineComment()) {
appendIfNotNull(result, sql.substring(startIndex, simpleParser.getPos()));
}
appendIfNotNull(result, sql.substring(startIndex));
return sql.length();
return simpleParser.getPos();
}

/** Skips a quoted string from startIndex. */
Expand Down
Loading

0 comments on commit 48828df

Please sign in to comment.