Skip to content

Commit 0daf8ed

Browse files
misoltmisolt
authored
Fix: Snowflake NPE on bad arguments (#968)
* Improve and test error handling * Update .gitignore * Simplify database selection logic * Test an exception message
1 parent 4ca89c5 commit 0daf8ed

File tree

3 files changed

+100
-49
lines changed

3 files changed

+100
-49
lines changed

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ node_modules
2222

2323
/dbsync/**/bin/
2424
/dumper/**/bin/
25+
/dumper/app/progress.log.*
2526
/permissions-migration/**/bin
2627

2728
# Gradle build folders

‎dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnector.java‎

Lines changed: 75 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import static com.google.common.base.CaseFormat.UPPER_CAMEL;
2020
import static com.google.common.base.CaseFormat.UPPER_UNDERSCORE;
21+
import static org.apache.hadoop.util.Preconditions.checkNotNull;
2122

2223
import com.google.common.base.CharMatcher;
2324
import com.google.common.base.Joiner;
@@ -41,9 +42,9 @@
4142
import java.util.ArrayList;
4243
import java.util.List;
4344
import java.util.Properties;
45+
import java.util.function.Supplier;
4446
import javax.annotation.Nonnull;
4547
import javax.sql.DataSource;
46-
import org.apache.commons.lang3.StringUtils;
4748
import org.springframework.jdbc.core.JdbcTemplate;
4849
import org.springframework.jdbc.datasource.SimpleDriverDataSource;
4950

@@ -72,31 +73,32 @@ public AbstractSnowflakeConnector(@Nonnull String name) {
7273
super(name);
7374
}
7475

75-
private static final int MAX_DATABASE_CHAR_LENGTH = 255;
76-
private static final String DEFAULT_DATABASE = "SNOWFLAKE";
77-
7876
@Nonnull
7977
@Override
8078
public abstract String getDescription();
8179

8280
@Nonnull
8381
@Override
84-
public Handle open(@Nonnull ConnectorArguments arguments)
82+
public final Handle open(@Nonnull ConnectorArguments arguments)
8583
throws MetadataDumperUsageException, SQLException {
86-
String url = arguments.getUri() != null ? arguments.getUri() : getUrlFromArguments(arguments);
87-
String databaseName =
88-
arguments.getDatabases().isEmpty()
89-
? DEFAULT_DATABASE
90-
: sanitizeDatabaseName(arguments.getDatabases().get(0));
91-
92-
DataSource dataSource =
93-
arguments.isPrivateKeyFileProvided()
94-
? createPrivateKeyDataSource(arguments, url)
95-
: createUserPasswordDataSource(arguments, url);
96-
JdbcHandle jdbcHandle = new JdbcHandle(dataSource);
97-
98-
setCurrentDatabase(databaseName, jdbcHandle.getJdbcTemplate());
99-
return jdbcHandle;
84+
Properties properties = dataSourceProperties(arguments);
85+
String url = getUrlFromArguments(arguments);
86+
DataSource dataSource = new SimpleDriverDataSource(newDriver(arguments), url, properties);
87+
if (arguments.isAssessment()) {
88+
JdbcHandle handle = new JdbcHandle(dataSource);
89+
JdbcTemplate template = handle.getJdbcTemplate();
90+
String actualDatabase = template.queryForObject("USE DATABASE SNOWFLAKE;", String.class);
91+
checkNotNull(actualDatabase);
92+
return handle;
93+
} else {
94+
String databaseName =
95+
arguments.getDatabases().isEmpty()
96+
? "SNOWFLAKE"
97+
: sanitizeDatabaseName(arguments.getDatabases().get(0));
98+
JdbcHandle handle = new JdbcHandle(dataSource);
99+
setCurrentDatabase(databaseName, handle.getJdbcTemplate());
100+
return handle;
101+
}
100102
}
101103

102104
@Override
@@ -121,40 +123,56 @@ public final void validate(@Nonnull ConnectorArguments arguments) {
121123
*/
122124
protected abstract void validateForConnector(@Nonnull ConnectorArguments arguments);
123125

124-
private DataSource createUserPasswordDataSource(@Nonnull ConnectorArguments arguments, String url)
126+
@Nonnull
127+
private Driver newDriver(@Nonnull ConnectorArguments arguments) throws SQLException {
128+
return newDriver(arguments.getDriverPaths(), "net.snowflake.client.jdbc.SnowflakeDriver");
129+
}
130+
131+
@Nonnull
132+
private static Properties dataSourceProperties(@Nonnull ConnectorArguments arguments)
125133
throws SQLException {
126-
Driver driver =
127-
newDriver(arguments.getDriverPaths(), "net.snowflake.client.jdbc.SnowflakeDriver");
128-
Properties prop = new Properties();
134+
String user = arguments.getUserOrFail();
135+
if (arguments.isPrivateKeyFileProvided()) {
136+
return createPrivateKeyProperties(arguments, user);
137+
} else {
138+
return createUserPasswordProperties(arguments, user);
139+
}
140+
}
141+
142+
private static Properties createUserPasswordProperties(
143+
@Nonnull ConnectorArguments arguments, @Nonnull String user) {
144+
Properties properties = new Properties();
129145

130-
prop.put("user", arguments.getUser());
146+
properties.put("user", user);
131147
if (arguments.isPasswordFlagProvided()) {
132-
prop.put("password", arguments.getPasswordOrPrompt());
148+
properties.put("password", arguments.getPasswordOrPrompt());
133149
}
134150
// Set default authenticator only if url is not provided to allow user overriding it
135151
if (arguments.getUri() == null) {
136-
prop.put("authenticator", "username_password_mfa");
152+
properties.put("authenticator", "username_password_mfa");
137153
}
138-
return new SimpleDriverDataSource(driver, url, prop);
154+
return properties;
139155
}
140156

141-
private DataSource createPrivateKeyDataSource(@Nonnull ConnectorArguments arguments, String url)
142-
throws SQLException {
143-
Driver driver =
144-
newDriver(arguments.getDriverPaths(), "net.snowflake.client.jdbc.SnowflakeDriver");
145-
Properties prop = new Properties();
157+
private static Properties createPrivateKeyProperties(
158+
@Nonnull ConnectorArguments arguments, @Nonnull String user) {
159+
Properties properties = new Properties();
160+
properties.put("user", user);
146161

147-
prop.put("private_key_file", arguments.getPrivateKeyFile());
148-
prop.put("user", arguments.getUser());
162+
properties.put("private_key_file", arguments.getPrivateKeyFile());
149163
if (arguments.getPrivateKeyPassword() != null) {
150-
prop.put("private_key_pwd", arguments.getPrivateKeyPassword());
164+
properties.put("private_key_pwd", arguments.getPrivateKeyPassword());
151165
}
152-
153-
return new SimpleDriverDataSource(driver, url, prop);
166+
return properties;
154167
}
155168

156169
@Nonnull
157170
private String getUrlFromArguments(@Nonnull ConnectorArguments arguments) {
171+
String url = arguments.getUri();
172+
if (url != null) {
173+
return url;
174+
}
175+
158176
StringBuilder buf = new StringBuilder("jdbc:snowflake://");
159177
String host = arguments.getHost("host.snowflakecomputing.com");
160178
buf.append(host).append("/");
@@ -178,26 +196,34 @@ private void setCurrentDatabase(@Nonnull String databaseName, @Nonnull JdbcTempl
178196
String currentDatabase =
179197
jdbcTemplate.queryForObject(String.format("USE DATABASE %s;", databaseName), String.class);
180198
if (currentDatabase == null) {
181-
List<String> dbNames =
182-
jdbcTemplate.query("SHOW DATABASES", (rs, rowNum) -> rs.getString("name"));
183-
throw new MetadataDumperUsageException(
184-
"Database name not found "
185-
+ databaseName
186-
+ ", use one of: "
187-
+ StringUtils.join(dbNames, ", "));
199+
Supplier<List<String>> showQuery =
200+
() -> jdbcTemplate.query("SHOW DATABASES", (rs, rowNum) -> rs.getString("name"));
201+
throw unrecognizedDatabase(databaseName, showQuery);
188202
}
189203
}
190204

205+
@Nonnull
206+
static MetadataDumperUsageException unrecognizedDatabase(
207+
@Nonnull String database, @Nonnull Supplier<List<String>> availableDatabases) {
208+
List<String> names = availableDatabases.get();
209+
String joinedNames = String.join(", ", names);
210+
String message =
211+
String.format("Database name not found %s, use one of: %s", database, joinedNames);
212+
213+
return new MetadataDumperUsageException(message);
214+
}
215+
191216
String sanitizeDatabaseName(@Nonnull String databaseName) throws MetadataDumperUsageException {
192-
CharMatcher doubleQuoteMatcher = CharMatcher.is('"');
193-
String trimmedName = doubleQuoteMatcher.trimFrom(databaseName);
194-
int charLengthWithQuotes = databaseName.length() + 2;
195-
if (charLengthWithQuotes > 255) {
217+
int lengthWithQuotes = databaseName.length() + 2;
218+
int maxLength = 255;
219+
if (lengthWithQuotes > maxLength) {
196220
throw new MetadataDumperUsageException(
197221
String.format(
198222
"The provided database name has %d characters, which is longer than the maximum allowed number %d for Snowflake identifiers.",
199-
charLengthWithQuotes, MAX_DATABASE_CHAR_LENGTH));
223+
lengthWithQuotes, maxLength));
200224
}
225+
CharMatcher doubleQuoteMatcher = CharMatcher.is('"');
226+
String trimmedName = doubleQuoteMatcher.trimFrom(databaseName);
201227
if (doubleQuoteMatcher.matchesAnyOf(trimmedName)) {
202228
throw new MetadataDumperUsageException(
203229
"Database name has incorrectly placed double quote(s). Aborting query.");

‎dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnectorTest.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
*/
1717
package com.google.edwmigration.dumper.application.dumper.connector.snowflake;
1818

19+
import static com.google.edwmigration.dumper.application.dumper.connector.snowflake.AbstractSnowflakeConnector.unrecognizedDatabase;
1920
import static org.junit.Assert.assertEquals;
21+
import static org.junit.Assert.assertNotNull;
2022
import static org.junit.Assert.assertThrows;
2123
import static org.junit.Assert.assertTrue;
2224

@@ -26,6 +28,8 @@
2628
import com.google.edwmigration.dumper.application.dumper.connector.AbstractConnectorTest;
2729
import java.io.IOException;
2830
import java.util.ArrayList;
31+
import java.util.List;
32+
import java.util.function.Supplier;
2933
import org.junit.Assert;
3034
import org.junit.Test;
3135
import org.junit.experimental.theories.Theories;
@@ -102,6 +106,14 @@ public void open_malformedInput_fail() throws IOException {
102106
e.getMessage().contains("Database name has incorrectly placed double quote(s)."));
103107
}
104108

109+
@Test
110+
public void open_noUser_throwsUsageException() throws Exception {
111+
ConnectorArguments arguments =
112+
ConnectorArguments.create(ImmutableList.of("--connector", "snowflake", "--assessment"));
113+
114+
assertThrows(MetadataDumperUsageException.class, () -> metadataConnector.open(arguments));
115+
}
116+
105117
@Test
106118
public void validate_mixedPrivateKeyAndPassword_fail() throws IOException {
107119
ConnectorArguments arguments =
@@ -119,6 +131,18 @@ public void validate_mixedPrivateKeyAndPassword_fail() throws IOException {
119131
"Private key authentication method can't be used together with user password"));
120132
}
121133

134+
@Test
135+
public void unrecognizedDatabase_success() {
136+
Supplier<List<String>> databases = () -> ImmutableList.of("SNOWFLAKE", "FIRSTDB", "SECONDDB");
137+
138+
String message = unrecognizedDatabase("WRONGNAMEDB", databases).getMessage();
139+
140+
assertNotNull(message);
141+
assertTrue(message, message.contains("WRONGNAMEDB"));
142+
assertTrue(message, message.contains("Database name not found"));
143+
assertTrue(message, message.contains("SNOWFLAKE, FIRSTDB, SECONDDB"));
144+
}
145+
122146
enum TestEnum {
123147
SomeValue
124148
}

0 commit comments

Comments
 (0)