Skip to content

Commit

Permalink
SNOW-1636286: Add exact search for schema
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-dprzybysz
sfc-gh-dprzybysz committed Dec 19, 2024
1 parent b3a2763 commit bf139a8
Show file tree
Hide file tree
Showing 6 changed files with 209 additions and 64 deletions.
13 changes: 13 additions & 0 deletions src/main/java/net/snowflake/client/core/SFBaseSession.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,11 @@ public abstract class SFBaseSession {
// we need to allow for it to maintain backwards compatibility.
private boolean enablePatternSearch = true;

// Enables the use of exact schema searches for certain DatabaseMetaData methods
// that should use schema from context (CLIENT_METADATA_REQUEST_USE_CONNECTION_CTX=true)
// value is false for backwards compatibility.
private boolean enableExactSchemaSearch = false;

/** Disable lookup for default credentials by GCS library */
private boolean disableGcsDefaultCredentials = false;

Expand Down Expand Up @@ -1069,6 +1074,14 @@ public void setEnablePatternSearch(boolean enablePatternSearch) {
this.enablePatternSearch = enablePatternSearch;
}

public boolean getEnableExactSchemaSearch() {
return enableExactSchemaSearch;
}

void setEnableExactSchemaSearch(boolean enableExactSchemaSearch) {
this.enableExactSchemaSearch = enableExactSchemaSearch;
}

public boolean getDisableGcsDefaultCredentials() {
return disableGcsDefaultCredentials;
}
Expand Down
7 changes: 7 additions & 0 deletions src/main/java/net/snowflake/client/core/SFSession.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -509,6 +509,13 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
setEnablePatternSearch(getBooleanValue(propertyValue));
}
break;

case ENABLE_EXACT_SCHEMA_SEARCH_ENABLED:
if (propertyValue != null) {
setEnableExactSchemaSearch(getBooleanValue(propertyValue));
}
break;

case DISABLE_GCS_DEFAULT_CREDENTIALS:
if (propertyValue != null) {
setDisableGcsDefaultCredentials(getBooleanValue(propertyValue));
Expand Down
1 change: 1 addition & 0 deletions src/main/java/net/snowflake/client/core/SFSessionProperty.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ public enum SFSessionProperty {
DIAGNOSTICS_ALLOWLIST_FILE("DIAGNOSTICS_ALLOWLIST_FILE", false, String.class),

ENABLE_PATTERN_SEARCH("enablePatternSearch", false, Boolean.class),
ENABLE_EXACT_SCHEMA_SEARCH_ENABLED("ENABLE_EXACT_SCHEMA_SEARCH_ENABLED", false, Boolean.class),

DISABLE_GCS_DEFAULT_CREDENTIALS("disableGcsDefaultCredentials", false, Boolean.class),

Expand Down
132 changes: 71 additions & 61 deletions src/main/java/net/snowflake/client/jdbc/SnowflakeDatabaseMetaData.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@
import net.snowflake.client.log.ArgSupplier;
import net.snowflake.client.log.SFLogger;
import net.snowflake.client.log.SFLoggerFactory;
import net.snowflake.client.util.SFPair;
import net.snowflake.client.util.SFTriple;
import net.snowflake.common.core.SqlState;
import net.snowflake.common.util.Wildcard;

Expand Down Expand Up @@ -167,6 +167,7 @@ public class SnowflakeDatabaseMetaData implements DatabaseMetaData {

// Indicates if pattern matching is allowed for all parameters.
private boolean isPatternMatchingEnabled = true;
private boolean exactSchemaSearchEnabled;

SnowflakeDatabaseMetaData(Connection connection) throws SQLException {
logger.trace("SnowflakeDatabaseMetaData(SnowflakeConnection connection)", false);
Expand All @@ -179,6 +180,7 @@ public class SnowflakeDatabaseMetaData implements DatabaseMetaData {
this.ibInstance = session.getTelemetryClient();
this.procedureResultsetColumnNum = -1;
this.isPatternMatchingEnabled = session.getEnablePatternSearch();
this.exactSchemaSearchEnabled = session.getEnableExactSchemaSearch();
}

private void raiseSQLExceptionIfConnectionIsClosed() throws SQLException {
Expand Down Expand Up @@ -1390,7 +1392,8 @@ else if (i == 0) {
}

// apply session context when catalog is unspecified
private SFPair<String, String> applySessionContext(String catalog, String schemaPattern) {
private SFTriple<String, String, Boolean> applySessionContext(
String catalog, String schemaPattern) {
if (metadataRequestUseConnectionCtx) {
// CLIENT_METADATA_USE_SESSION_DATABASE = TRUE
if (catalog == null) {
Expand All @@ -1407,17 +1410,18 @@ private SFPair<String, String> applySessionContext(String catalog, String schema
}
}
}
return SFPair.of(catalog, schemaPattern);
return new SFTriple<>(catalog, schemaPattern, exactSchemaSearchEnabled && useSessionSchema);
}

/* helper function for getProcedures, getFunctionColumns, etc. Returns sql command to show some type of result such
as procedures or udfs */
private String getFirstResultSetCommand(
String catalog, String schemaPattern, String name, String type) {
// apply session context when catalog is unspecified
SFPair<String, String> resPair = applySessionContext(catalog, schemaPattern);
catalog = resPair.left;
schemaPattern = resPair.right;
SFTriple<String, String, Boolean> result = applySessionContext(catalog, schemaPattern);
catalog = result.first();
schemaPattern = result.second();
boolean isExactSchema = result.third();

String showProcedureCommand = "show /* JDBC:DatabaseMetaData.getProcedures() */ " + type;

Expand All @@ -1431,7 +1435,7 @@ private String getFirstResultSetCommand(
return "";
} else {
String catalogEscaped = escapeSqlQuotes(catalog);
if (schemaPattern == null || isSchemaNameWildcardPattern(schemaPattern)) {
if (!isExactSchema && (schemaPattern == null || isSchemaNameWildcardPattern(schemaPattern))) {
showProcedureCommand += " in database \"" + catalogEscaped + "\"";
} else if (schemaPattern.isEmpty()) {
return "";
Expand Down Expand Up @@ -1510,9 +1514,11 @@ public ResultSet getTables(
return SnowflakeDatabaseMetaDataResultSet.getEmptyResultSet(GET_TABLES, statement);
}

SFPair<String, String> resPair = applySessionContext(originalCatalog, originalSchemaPattern);
final String catalog = resPair.left;
final String schemaPattern = resPair.right;
SFTriple<String, String, Boolean> result =
applySessionContext(originalCatalog, originalSchemaPattern);
String catalog = result.first();
String schemaPattern = result.second();
boolean isExactSchema = result.third();

final Pattern compiledSchemaPattern = Wildcard.toRegexPattern(schemaPattern, true);
final Pattern compiledTablePattern = Wildcard.toRegexPattern(tableNamePattern, true);
Expand Down Expand Up @@ -1553,7 +1559,7 @@ public ResultSet getTables(
} else if (schemaPattern.isEmpty()) {
return SnowflakeDatabaseMetaDataResultSet.getEmptyResultSet(GET_TABLES, statement);
} else {
String schemaUnescaped = unescapeChars(schemaPattern);
String schemaUnescaped = isExactSchema ? schemaPattern : unescapeChars(schemaPattern);
showTablesCommand += " in schema \"" + catalogEscaped + "\".\"" + schemaUnescaped + "\"";
}
}
Expand Down Expand Up @@ -1699,9 +1705,11 @@ public ResultSet getColumns(
Statement statement = connection.createStatement();

// apply session context when catalog is unspecified
SFPair<String, String> resPair = applySessionContext(originalCatalog, originalSchemaPattern);
final String catalog = resPair.left;
final String schemaPattern = resPair.right;
SFTriple<String, String, Boolean> result =
applySessionContext(originalCatalog, originalSchemaPattern);
String catalog = result.first();
String schemaPattern = result.second();
boolean isExactSchema = result.third();

final Pattern compiledSchemaPattern = Wildcard.toRegexPattern(schemaPattern, true);
final Pattern compiledTablePattern = Wildcard.toRegexPattern(tableNamePattern, true);
Expand Down Expand Up @@ -1729,7 +1737,7 @@ public ResultSet getColumns(
return SnowflakeDatabaseMetaDataResultSet.getEmptyResultSet(
extendedSet ? GET_COLUMNS_EXTENDED_SET : GET_COLUMNS, statement);
} else {
String schemaUnescaped = unescapeChars(schemaPattern);
String schemaUnescaped = isExactSchema ? schemaPattern : unescapeChars(schemaPattern);
if (tableNamePattern == null || Wildcard.isWildcardPatternStr(tableNamePattern)) {
showColumnsCommand += " in schema \"" + catalogEscaped + "\".\"" + schemaUnescaped + "\"";
} else if (tableNamePattern.isEmpty()) {
Expand Down Expand Up @@ -1968,9 +1976,11 @@ public ResultSet getTablePrivileges(
return SnowflakeDatabaseMetaDataResultSet.getEmptyResultSet(GET_TABLE_PRIVILEGES, statement);
}
// apply session context when catalog is unspecified
SFPair<String, String> resPair = applySessionContext(originalCatalog, originalSchemaPattern);
final String catalog = resPair.left;
final String schemaPattern = resPair.right;
SFTriple<String, String, Boolean> result =
applySessionContext(originalCatalog, originalSchemaPattern);
String catalog = result.first();
String schemaPattern = result.second();
boolean isExactSchema = result.third();

String showView = "select * from ";

Expand All @@ -1993,10 +2003,11 @@ public ResultSet getTablePrivileges(
&& !schemaPattern.isEmpty()
&& !schemaPattern.trim().equals("%")
&& !schemaPattern.trim().equals(".*")) {
String unescapedSchema = isExactSchema ? schemaPattern : unescapeChars(schemaPattern);
if (showView.contains("where table_name")) {
showView += " and table_schema = '" + unescapeChars(schemaPattern) + "'";
showView += " and table_schema = '" + unescapedSchema + "'";
} else {
showView += " where table_schema = '" + unescapeChars(schemaPattern) + "'";
showView += " where table_schema = '" + unescapedSchema + "'";
}
}
showView += " order by table_catalog, table_schema, table_name, privilege_type";
Expand Down Expand Up @@ -2087,9 +2098,10 @@ public ResultSet getPrimaryKeys(String originalCatalog, String originalSchema, f
String showPKCommand = "show /* JDBC:DatabaseMetaData.getPrimaryKeys() */ primary keys in ";

// apply session context when catalog is unspecified
SFPair<String, String> resPair = applySessionContext(originalCatalog, originalSchema);
final String catalog = resPair.left;
final String schema = resPair.right;
SFTriple<String, String, Boolean> result = applySessionContext(originalCatalog, originalSchema);
String catalog = result.first();
String schema = result.second();
boolean isExactSchema = result.third();

// These Patterns will only be used if the connection property enablePatternSearch=true
final Pattern compiledSchemaPattern = Wildcard.toRegexPattern(schema, true);
Expand All @@ -2106,7 +2118,7 @@ public ResultSet getPrimaryKeys(String originalCatalog, String originalSchema, f
} else if (schema.isEmpty()) {
return SnowflakeDatabaseMetaDataResultSet.getEmptyResultSet(GET_PRIMARY_KEYS, statement);
} else {
String schemaUnescaped = unescapeChars(schema);
String schemaUnescaped = isExactSchema ? schema : unescapeChars(schema);
if (table == null) {
showPKCommand += "schema \"" + catalogUnescaped + "\".\"" + schemaUnescaped + "\"";
} else if (table.isEmpty()) {
Expand Down Expand Up @@ -2250,11 +2262,11 @@ private ResultSet getForeignKeys(
StringBuilder commandBuilder = new StringBuilder();

// apply session context when catalog is unspecified
// apply session context when catalog is unspecified
SFPair<String, String> resPair =
SFTriple<String, String, Boolean> result =
applySessionContext(originalParentCatalog, originalParentSchema);
final String parentCatalog = resPair.left;
final String parentSchema = resPair.right;
String parentCatalog = result.first();
String parentSchema = result.second();
boolean isExactSchema = result.third();

// These Patterns will only be used to filter results if the connection property
// enablePatternSearch=true
Expand Down Expand Up @@ -2283,7 +2295,7 @@ private ResultSet getForeignKeys(
} else if (parentSchema.isEmpty()) {
return SnowflakeDatabaseMetaDataResultSet.getEmptyResultSet(GET_FOREIGN_KEYS, statement);
} else {
String unescapedParentSchema = unescapeChars(parentSchema);
String unescapedParentSchema = isExactSchema ? parentSchema : unescapeChars(parentSchema);
if (parentTable == null) {
commandBuilder.append(
"schema \"" + unescapedParentCatalog + "\".\"" + unescapedParentSchema + "\"");
Expand Down Expand Up @@ -2580,12 +2592,7 @@ public ResultSet getImportedKeys(String originalCatalog, String originalSchema,
originalSchema,
table);

// apply session context when catalog is unspecified
SFPair<String, String> resPair = applySessionContext(originalCatalog, originalSchema);
final String catalog = resPair.left;
final String schema = resPair.right;

return getForeignKeys("import", catalog, schema, table, null, null, null);
return getForeignKeys("import", originalCatalog, originalSchema, table, null, null, null);
}

@Override
Expand All @@ -2598,10 +2605,6 @@ public ResultSet getExportedKeys(String catalog, String schema, String table)
schema,
table);

// apply session context when catalog is unspecified
SFPair<String, String> resPair = applySessionContext(catalog, schema);
catalog = resPair.left;
schema = resPair.right;
return getForeignKeys("export", catalog, schema, table, null, null, null);
}

Expand All @@ -2625,11 +2628,6 @@ public ResultSet getCrossReference(
foreignCatalog,
foreignSchema,
foreignTable);
// apply session context when catalog is unspecified
SFPair<String, String> resPair = applySessionContext(parentCatalog, parentSchema);
parentCatalog = resPair.left;
parentSchema = resPair.right;

return getForeignKeys(
"cross",
parentCatalog,
Expand Down Expand Up @@ -2877,9 +2875,11 @@ public ResultSet getStreams(
Statement statement = connection.createStatement();

// apply session context when catalog is unspecified
SFPair<String, String> resPair = applySessionContext(originalCatalog, originalSchemaPattern);
final String catalog = resPair.left;
final String schemaPattern = resPair.right;
SFTriple<String, String, Boolean> result =
applySessionContext(originalCatalog, originalSchemaPattern);
String catalog = result.first();
String schemaPattern = result.second();
boolean isExactSchema = result.third();

final Pattern compiledSchemaPattern = Wildcard.toRegexPattern(schemaPattern, true);
final Pattern compiledStreamNamePattern = Wildcard.toRegexPattern(streamName, true);
Expand All @@ -2904,7 +2904,7 @@ public ResultSet getStreams(
} else if (schemaPattern.isEmpty()) {
return SnowflakeDatabaseMetaDataResultSet.getEmptyResultSet(GET_STREAMS, statement);
} else {
String schemaUnescaped = unescapeChars(schemaPattern);
String schemaUnescaped = isExactSchema ? schemaPattern : unescapeChars(schemaPattern);
if (streamName == null || Wildcard.isWildcardPatternStr(streamName)) {
showStreamsCommand += " in schema \"" + catalogEscaped + "\".\"" + schemaUnescaped + "\"";
}
Expand Down Expand Up @@ -3280,35 +3280,44 @@ public ResultSet getSchemas(String originalCatalog, String originalSchema) throw
raiseSQLExceptionIfConnectionIsClosed();

// apply session context when catalog is unspecified
SFPair<String, String> resPair = applySessionContext(originalCatalog, originalSchema);
final String catalog = resPair.left;
final String schemaPattern = resPair.right;
SFTriple<String, String, Boolean> result = applySessionContext(originalCatalog, originalSchema);
final String catalog = result.first();
final String schemaPattern = result.second();
boolean isExactSchema = result.third();

final Pattern compiledSchemaPattern = Wildcard.toRegexPattern(schemaPattern, true);

String showSchemas = "show /* JDBC:DatabaseMetaData.getSchemas() */ schemas";
StringBuilder showSchemas =
new StringBuilder("show /* JDBC:DatabaseMetaData.getSchemas() */ schemas");

Statement statement = connection.createStatement();
// only add pattern if it is not empty and not matching all character.
if (schemaPattern != null
if (isExactSchema) {
String escapedSchema =
schemaPattern.replaceAll("_", "\\\\\\\\_").replaceAll("%", "\\\\\\\\%");
showSchemas.append(" like '").append(escapedSchema).append("'");
} else if (schemaPattern != null
&& !schemaPattern.isEmpty()
&& !schemaPattern.trim().equals("%")
&& !schemaPattern.trim().equals(".*")) {
showSchemas += " like '" + escapeSingleQuoteForLikeCommand(schemaPattern) + "'";
// only add pattern if it is not empty and not matching all character.
showSchemas
.append(" like '")
.append(escapeSingleQuoteForLikeCommand(schemaPattern))
.append("'");
}

if (catalog == null) {
showSchemas += " in account";
showSchemas.append(" in account");
} else if (catalog.isEmpty()) {
return SnowflakeDatabaseMetaDataResultSet.getEmptyResultSet(GET_SCHEMAS, statement);
} else {
showSchemas += " in database \"" + escapeSqlQuotes(catalog) + "\"";
showSchemas.append(" in database \"").append(escapeSqlQuotes(catalog)).append("\"");
}

logger.debug("Sql command to get schemas metadata: {}", showSchemas);
String sqlQuery = showSchemas.toString();
logger.debug("Sql command to get schemas metadata: {}", sqlQuery);

ResultSet resultSet =
executeAndReturnEmptyResultIfNotFound(statement, showSchemas, GET_SCHEMAS);
ResultSet resultSet = executeAndReturnEmptyResultIfNotFound(statement, sqlQuery, GET_SCHEMAS);
sendInBandTelemetryMetadataMetrics(
resultSet, "getSchemas", originalCatalog, originalSchema, "none", "none");
return new SnowflakeDatabaseMetaDataQueryResultSet(GET_SCHEMAS, resultSet, statement) {
Expand All @@ -3323,7 +3332,8 @@ public boolean next() throws SQLException {
String dbName = showObjectResultSet.getString(5);

if (compiledSchemaPattern == null
|| compiledSchemaPattern.matcher(schemaName).matches()) {
|| compiledSchemaPattern.matcher(schemaName).matches()
|| isExactSchema) {
nextRow[0] = schemaName;
nextRow[1] = dbName;
return true;
Expand Down
Loading

0 comments on commit bf139a8

Please sign in to comment.