Skip to content

Commit 008e9e1

Browse files
Bruce Irschickbirschick-bq
andauthored
[AD-974] Allow the option allowDiskUse to be set via an environment variable. (#422)
* [AD-890] Ensure JDBC driver supports and tests Decimal128 type * Commit Code Coverage Badge * [AD-890] Code review comments about DECIMAL literals. * [AD-890] Add improved DECIMAL call to server. * [AD-922] User can enable/disable AllowDiskUse query option for a connection. * Commit Code Coverage Badge * [AD-922] Correct problem with Tableau connector configuration. * Commit Code Coverage Badge * [AD-922] Change default to DEFAULT and ensure propagation from connection properties. * Commit Code Coverage Badge * [AD-974] Revert changes to Tableau connector. * [AD-974] Enable allowDiskUse option to be set via an environment variable DOCUMENTDB_CUSTOM_OPTIONS. Currently, undocumented. * Commit Code Coverage Badge * [AD-974] Fix code review comments. * Commit Code Coverage Badge Co-authored-by: birschick-bq <birschick-bq@users.noreply.github.com>
1 parent c6b70cb commit 008e9e1

File tree

9 files changed

+321
-41
lines changed

9 files changed

+321
-41
lines changed
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
* Copyright <2021> Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*
15+
*/
16+
17+
package software.amazon.documentdb.jdbc;
18+
19+
/**
20+
* The enumeration of Allow Disk Use options.
21+
*/
22+
public enum DocumentDbAllowDiskUseOption {
23+
DEFAULT("default"),
24+
DISABLE("disable"),
25+
ENABLE("enable"),
26+
;
27+
28+
private final String name;
29+
30+
DocumentDbAllowDiskUseOption(final String name) {
31+
this.name = name;
32+
}
33+
34+
public String getName() {
35+
return name;
36+
}
37+
38+
/**
39+
* Returns DocumentDbAllowDiskUseOption with a name that matches input string.
40+
* @param allowDiskUseOption name of the allow disk use option.
41+
* @return DocumentDbAllowDiskUseOption of string.
42+
*/
43+
public static DocumentDbAllowDiskUseOption fromString(final String allowDiskUseOption) {
44+
for (DocumentDbAllowDiskUseOption scanMethod: DocumentDbAllowDiskUseOption.values()) {
45+
if (scanMethod.name.equals(allowDiskUseOption)) {
46+
return scanMethod;
47+
}
48+
}
49+
throw new IllegalArgumentException("Invalid allow disk use option.");
50+
}
51+
}

src/main/java/software/amazon/documentdb/jdbc/DocumentDbConnectionProperties.java

Lines changed: 86 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ public class DocumentDbConnectionProperties extends Properties {
6565
private static final String ROOT_PEM_RESOURCE_FILE_NAME = "/rds-ca-2019-root.pem";
6666
public static final String HOME_PATH_PREFIX_REG_EXPR = "^~[/\\\\].*$";
6767
public static final int FETCH_SIZE_DEFAULT = 2000;
68+
public static final String DOCUMENTDB_CUSTOM_OPTIONS = "DOCUMENTDB_CUSTOM_OPTIONS";
6869
private static String classPathLocationName = null;
6970
private static String[] sshPrivateKeyFileSearchPaths = null;
7071
private static final String DEFAULT_APPLICATION_NAME_KEY = "default.application.name";
@@ -482,7 +483,7 @@ public String getSshUser() {
482483
}
483484

484485
/**
485-
* Sets the SSH tunnel host name. Can optionally contain the port number using `host-name:port'
486+
* Sets the SSH tunnel host name. Can optionally contain the port number using 'host-name:port'
486487
* syntax. If port is not provided, port 22 is assumed.
487488
*
488489
* @param sshHostname the SSH tunnel host name and optional port number.
@@ -658,6 +659,24 @@ public String getDefaultAuthenticationDatabase() {
658659
DocumentDbConnectionProperty.DEFAULT_AUTH_DB.getDefaultValue());
659660
}
660661

662+
/**
663+
* Sets the allow disk use option.
664+
*
665+
* @param allowDiskUseOption the disk use option to set.
666+
*/
667+
public void setAllowDiskUseOption(final String allowDiskUseOption) {
668+
setProperty(DocumentDbConnectionProperty.ALLOW_DISK_USE.getName(), allowDiskUseOption);
669+
}
670+
671+
/**
672+
* Gets the allow disk use option.
673+
*
674+
* @return the disk use option, or null, if invalid or not set.
675+
*/
676+
public DocumentDbAllowDiskUseOption getAllowDiskUseOption() {
677+
return getPropertyAsAllowDiskUseOption(DocumentDbConnectionProperty.ALLOW_DISK_USE.getName());
678+
}
679+
661680
/**
662681
* Builds the MongoClientSettings from properties.
663682
*
@@ -813,6 +832,9 @@ public String buildSanitizedConnectionString() {
813832
if (getDefaultAuthenticationDatabase() != null && !DocumentDbConnectionProperty.DEFAULT_AUTH_DB.getDefaultValue().equals(getDefaultAuthenticationDatabase())) {
814833
appendOption(optionalInfo, DocumentDbConnectionProperty.DEFAULT_AUTH_DB, getDefaultAuthenticationDatabase());
815834
}
835+
if (getAllowDiskUseOption() != DocumentDbAllowDiskUseOption.fromString(DocumentDbConnectionProperty.ALLOW_DISK_USE.getDefaultValue())) {
836+
appendOption(optionalInfo, DocumentDbConnectionProperty.ALLOW_DISK_USE, getAllowDiskUseOption().getName());
837+
}
816838
return String.format(connectionStringTemplate,
817839
loginInfo,
818840
hostInfo,
@@ -896,6 +918,8 @@ public static DocumentDbConnectionProperties getPropertiesFromConnectionString(
896918

897919
setOptionalProperties(properties, uri);
898920

921+
setCustomOptions(properties);
922+
899923
} catch (URISyntaxException e) {
900924
throw SqlError.createSQLException(
901925
LOGGER,
@@ -912,6 +936,21 @@ public static DocumentDbConnectionProperties getPropertiesFromConnectionString(
912936
return properties;
913937
}
914938

939+
static void setCustomOptions(final DocumentDbConnectionProperties properties) {
940+
final String customOptions = System.getenv(DOCUMENTDB_CUSTOM_OPTIONS);
941+
if (customOptions == null) {
942+
return;
943+
}
944+
final String[] propertyPairs = customOptions.split(";");
945+
for (String pair : propertyPairs) {
946+
final int splitIndex = pair.indexOf("=");
947+
final String key = pair.substring(0, splitIndex);
948+
final String value = pair.substring(1 + splitIndex);
949+
950+
addPropertyIfValid(properties, key, value, true, true);
951+
}
952+
}
953+
915954
private static void setDatabase(final Properties properties, final URI mongoUri)
916955
throws SQLException {
917956
if (isNullOrWhitespace(mongoUri.getPath())) {
@@ -942,7 +981,7 @@ private static void setOptionalProperties(final Properties properties, final URI
942981
final String key = pair.substring(0, splitIndex);
943982
final String value = pair.substring(1 + splitIndex);
944983

945-
addPropertyIfValid(properties, key, value);
984+
addPropertyIfValid(properties, key, value, false, false);
946985
}
947986
}
948987

@@ -997,15 +1036,33 @@ private static void setHostName(final Properties properties, final URI mongoUri)
9971036
}
9981037

9991038
private static void addPropertyIfValid(
1000-
final Properties info, final String propertyKey, final String propertyValue) {
1039+
final Properties info,
1040+
final String propertyKey,
1041+
final String propertyValue,
1042+
final boolean allowUnsupported,
1043+
final boolean allowUnknown) {
10011044
if (DocumentDbConnectionProperty.isSupportedProperty(propertyKey)) {
10021045
addPropertyIfNotSet(info, propertyKey, propertyValue);
10031046
} else if (DocumentDbConnectionProperty.isUnsupportedMongoDBProperty(propertyKey)) {
1004-
LOGGER.warn(
1005-
"Ignored MongoDB property: {{}} as it not supported by the driver.",
1006-
propertyKey);
1047+
if (allowUnsupported) {
1048+
LOGGER.warn(
1049+
"Adding unsupported MongoDB property: {{}} it may not supported by the driver or server.",
1050+
propertyKey);
1051+
addPropertyIfNotSet(info, propertyKey, propertyValue);
1052+
} else {
1053+
LOGGER.warn(
1054+
"Ignored MongoDB property: {{}} as it not supported by the driver.",
1055+
propertyKey);
1056+
}
10071057
} else {
1008-
LOGGER.warn("Ignored invalid property: {{}}", propertyKey);
1058+
if (allowUnknown) {
1059+
LOGGER.warn(
1060+
"Adding unknown MongoDB property: {{}} it may not supported by the driver or server.",
1061+
propertyKey);
1062+
addPropertyIfNotSet(info, propertyKey, propertyValue);
1063+
} else {
1064+
LOGGER.warn("Ignored invalid property: {{}}", propertyKey);
1065+
}
10091066
}
10101067
}
10111068

@@ -1224,6 +1281,28 @@ private DocumentDbMetadataScanMethod getPropertyAsScanMethod(@NonNull final Stri
12241281
return property;
12251282
}
12261283

1284+
/**
1285+
* Attempts to retrieve a property as a DocumentDbAllowDiskUseOption.
1286+
*
1287+
* @param key The property to retrieve.
1288+
* @return The retrieved property as a DocumentDbAllowDiskUseOption or null if it did not exist or was not a
1289+
* valid DocumentDbAllowDiskUseOption.
1290+
*/
1291+
private DocumentDbAllowDiskUseOption getPropertyAsAllowDiskUseOption(@NonNull final String key) {
1292+
DocumentDbAllowDiskUseOption property = null;
1293+
try {
1294+
if (getProperty(key) != null) {
1295+
property = DocumentDbAllowDiskUseOption.fromString(getProperty(key));
1296+
} else if (DocumentDbConnectionProperty.getPropertyFromKey(key) != null) {
1297+
property = DocumentDbAllowDiskUseOption.fromString(
1298+
DocumentDbConnectionProperty.getPropertyFromKey(key).getDefaultValue());
1299+
}
1300+
} catch (IllegalArgumentException e) {
1301+
LOGGER.warn("Property {{}} was ignored as it was not a valid allow disk use option.", key, e);
1302+
}
1303+
return property;
1304+
}
1305+
12271306
/**
12281307
* Attempts to retrieve a property as a Long.
12291308
*

src/main/java/software/amazon/documentdb/jdbc/DocumentDbConnectionProperty.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,8 @@ public enum DocumentDbConnectionProperty implements ConnectionProperty {
7373
REFRESH_SCHEMA("refreshSchema", "false",
7474
"Refreshes any existing schema with a newly generated schema when the connection first requires the schema. Note that this will remove any existing schema customizations and will reduce performance for the first query or metadata inquiry."),
7575
DEFAULT_AUTH_DB("defaultAuthDb", "admin", "The default authentication database to use."),
76+
ALLOW_DISK_USE("allowDiskUse", "default",
77+
"Indicator of whether to enable or disable the 'allow disk use' option on all queries. Valid values are 'enable', 'disable' or 'default'. Default is 'default'."),
7678
;
7779

7880
// Unsupported MongoDB connection properties that will be ignored but should have warnings.

src/main/java/software/amazon/documentdb/jdbc/DocumentDbPreparedStatement.java

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
public class DocumentDbPreparedStatement extends PreparedStatement
3434
implements java.sql.PreparedStatement {
3535
private int queryTimeout = 0;
36+
private DocumentDbAllowDiskUseOption allowDiskUse = DocumentDbAllowDiskUseOption.DEFAULT;
3637
private final DocumentDbQueryExecutor queryExecutor;
3738

3839
/**
@@ -50,6 +51,7 @@ public DocumentDbPreparedStatement(final Connection connection, final String sql
5051
final DocumentDbQueryMappingService mappingService = new DocumentDbQueryMappingService(
5152
connectionProperties,
5253
documentDbConnection.getDatabaseMetadata());
54+
setAllowDiskUse(connectionProperties.getAllowDiskUseOption());
5355
queryExecutor = new DocumentDbQueryExecutor(
5456
this,
5557
connectionProperties,
@@ -108,4 +110,26 @@ public void setQueryTimeout(final int seconds) throws SQLException {
108110
queryTimeout = seconds;
109111
queryExecutor.setQueryTimeout(seconds);
110112
}
113+
114+
/**
115+
* Gets the allow disk use option for the statement.
116+
*
117+
* @return one of the allow disk use options.
118+
* @throws SQLException if the connection is not open.
119+
*/
120+
public DocumentDbAllowDiskUseOption getAllowDiskUse() throws SQLException {
121+
verifyOpen();
122+
return allowDiskUse;
123+
}
124+
125+
/**
126+
* Sets the allow disk use indicator for the statement.
127+
*
128+
* @param allowDiskUse the indicator of whether to set the allow disk use option.
129+
* @throws SQLException if the connection is not open.
130+
*/
131+
public void setAllowDiskUse(final DocumentDbAllowDiskUseOption allowDiskUse) throws SQLException {
132+
verifyOpen();
133+
this.allowDiskUse = allowDiskUse;
134+
}
111135
}

src/main/java/software/amazon/documentdb/jdbc/DocumentDbQueryExecutor.java

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ public class DocumentDbQueryExecutor {
5454
private final DocumentDbQueryMappingService queryMapper;
5555
private int fetchSize;
5656
private int queryTimeout;
57+
private DocumentDbAllowDiskUseOption allowDiskUse;
5758
private String queryId = null;
5859
private QueryState queryState = QueryState.NOT_STARTED;
5960

@@ -77,6 +78,7 @@ private enum QueryState {
7778
this.queryMapper = queryMapper;
7879
this.fetchSize = fetchSize;
7980
this.queryTimeout = queryTimeoutSecs;
81+
this.allowDiskUse = connectionProperties.getAllowDiskUseOption();
8082
}
8183

8284
/**
@@ -193,6 +195,11 @@ protected java.sql.ResultSet runQuery(final String sql) throws SQLException {
193195
if (getFetchSize() > 0) {
194196
iterable = iterable.batchSize(getFetchSize());
195197
}
198+
if (getAllowDiskUse() == DocumentDbAllowDiskUseOption.ENABLE) {
199+
iterable = iterable.allowDiskUse(true);
200+
} else if (getAllowDiskUse() == DocumentDbAllowDiskUseOption.DISABLE) {
201+
iterable = iterable.allowDiskUse(false);
202+
}
196203

197204
final ImmutableList<JdbcColumnMetaData> columnMetaData = ImmutableList
198205
.copyOf(queryContext.getColumnMetaData());
@@ -301,4 +308,13 @@ protected int getFetchSize() {
301308
protected void setFetchSize(final int fetchSize) {
302309
this.fetchSize = fetchSize;
303310
}
311+
312+
protected DocumentDbAllowDiskUseOption getAllowDiskUse() {
313+
return allowDiskUse;
314+
}
315+
316+
protected void setAllowDiskUse(final DocumentDbAllowDiskUseOption allowDiskUse) {
317+
this.allowDiskUse = allowDiskUse;
318+
}
319+
304320
}

0 commit comments

Comments
 (0)