Skip to content

feat: expose some driver settings #3739

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
fbiville wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat: expose some driver settings #3739

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions ...d-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/database/Neo4jConnection.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import java.io.Serializable;
import java.util.Collections;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;
import org.apache.commons.lang3.StringUtils;
import org.neo4j.driver.Config;
Expand Down Expand Up @@ -54,6 +55,10 @@ public Neo4jConnection(ConnectionParams settings, String templateVersion) {
settings.asAuthToken(),
Config.builder()
.withLogging(Logging.slf4j())
.withMaxConnectionLifetime(
maxConnectionLifetimeMillis(settings), TimeUnit.MILLISECONDS)
.withConnectionLivenessCheckTimeout(
connectionLivenessCheckTimeoutMillis(settings), TimeUnit.MILLISECONDS)
.withUserAgent(Neo4jTelemetry.userAgent(templateVersion))
.build()));
}
Expand Down Expand Up @@ -239,4 +244,20 @@ private static TransactionConfig databaseResetMetadata(String resetMethod) {
Neo4jTelemetry.transactionMetadata(Map.of("sink", "neo4j", "step", resetMethod)))
.build();
}

private static long maxConnectionLifetimeMillis(ConnectionParams settings) {
Long maxConnectionLifetimeMillis = settings.getMaxConnectionLifetimeMillis();
if (maxConnectionLifetimeMillis == null) {
return Config.defaultConfig().maxConnectionLifetimeMillis();
}
return maxConnectionLifetimeMillis;
}

private static Long connectionLivenessCheckTimeoutMillis(ConnectionParams settings) {
Long connectionLivenessCheckTimeoutMillis = settings.getConnectionLivenessCheckTimeoutMillis();
if (connectionLivenessCheckTimeoutMillis == null) {
return Config.defaultConfig().idleTimeBeforeConnectionTest();
}
return connectionLivenessCheckTimeoutMillis;
}
}
7 changes: 5 additions & 2 deletions .../main/java/com/google/cloud/teleport/v2/neo4j/model/connection/BasicConnectionParams.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,11 @@ public BasicConnectionParams(
@JsonProperty("server_url") String serverUrl,
@JsonProperty("database") String database,
@JsonProperty("username") String username,
@JsonProperty("pwd") String password) {
super(serverUrl, database);
@JsonProperty("pwd") String password,
@JsonProperty("max_connection_lifetime_millis") Long maxConnectionLifetimeMillis,
@JsonProperty("connection_liveness_check_timeout_millis")
Long connectionLivenessCheckTimeoutMillis) {
super(serverUrl, database, maxConnectionLifetimeMillis, connectionLivenessCheckTimeoutMillis);
this.username = username;
this.password = password;
}
Expand Down
7 changes: 5 additions & 2 deletions ...main/java/com/google/cloud/teleport/v2/neo4j/model/connection/BearerConnectionParams.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,11 @@ class BearerConnectionParams extends ConnectionParams {
public BearerConnectionParams(
@JsonProperty("server_url") String serverUrl,
@JsonProperty("database") String database,
@JsonProperty("token") String token) {
super(serverUrl, database);
@JsonProperty("token") String token,
@JsonProperty("max_connection_lifetime_millis") Long maxConnectionLifetimeMillis,
@JsonProperty("connection_liveness_check_timeout_millis")
Long connectionLivenessCheckTimeoutMillis) {
super(serverUrl, database, maxConnectionLifetimeMillis, connectionLivenessCheckTimeoutMillis);
this.token = token;
}

Expand Down
34 changes: 25 additions & 9 deletions ...j/src/main/java/com/google/cloud/teleport/v2/neo4j/model/connection/ConnectionParams.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,10 +34,18 @@ public abstract class ConnectionParams implements Serializable {

private final String serverUrl;
private final String database;
private final Long maxConnectionLifetimeMillis;
private final Long connectionLivenessCheckTimeoutMillis;

public ConnectionParams(String serverUrl, String database) {
public ConnectionParams(
String serverUrl,
String database,
Long maxConnectionLifetimeMillis,
Long connectionLivenessCheckTimeoutMillis) {
this.serverUrl = serverUrl;
this.database = database == null ? "neo4j" : database;
this.maxConnectionLifetimeMillis = maxConnectionLifetimeMillis;
this.connectionLivenessCheckTimeoutMillis = connectionLivenessCheckTimeoutMillis;
}

public String getServerUrl() {
Expand All @@ -48,23 +56,31 @@ public String getDatabase() {
return database;
}

public Long getMaxConnectionLifetimeMillis() {
return maxConnectionLifetimeMillis;
}

public Long getConnectionLivenessCheckTimeoutMillis() {
return connectionLivenessCheckTimeoutMillis;
}

public abstract AuthToken asAuthToken();

@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
if (!(o instanceof ConnectionParams that)) {
return false;
}

ConnectionParams that = (ConnectionParams) o;
return Objects.equals(serverUrl, that.serverUrl) && Objects.equals(database, that.database);
return Objects.equals(serverUrl, that.serverUrl)
&& Objects.equals(database, that.database)
&& Objects.equals(maxConnectionLifetimeMillis, that.maxConnectionLifetimeMillis)
&& Objects.equals(
connectionLivenessCheckTimeoutMillis, that.connectionLivenessCheckTimeoutMillis);
}

@Override
public int hashCode() {
return Objects.hash(serverUrl, database);
return Objects.hash(
serverUrl, database, maxConnectionLifetimeMillis, connectionLivenessCheckTimeoutMillis);
}
}
7 changes: 5 additions & 2 deletions ...main/java/com/google/cloud/teleport/v2/neo4j/model/connection/CustomConnectionParams.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,11 @@ public CustomConnectionParams(
@JsonProperty("credentials") String credentials,
@JsonProperty("realm") String realm,
@JsonProperty("scheme") String scheme,
@JsonProperty("parameters") Map<String, Object> parameters) {
super(serverUrl, database);
@JsonProperty("parameters") Map<String, Object> parameters,
@JsonProperty("max_connection_lifetime_millis") Long maxConnectionLifetimeMillis,
@JsonProperty("connection_liveness_check_timeout_millis")
Long connectionLivenessCheckTimeoutMillis) {
super(serverUrl, database, maxConnectionLifetimeMillis, connectionLivenessCheckTimeoutMillis);
this.principal = principal;
this.credentials = credentials;
this.realm = realm;
Expand Down
7 changes: 5 additions & 2 deletions ...in/java/com/google/cloud/teleport/v2/neo4j/model/connection/KerberosConnectionParams.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,11 @@ class KerberosConnectionParams extends ConnectionParams {
public KerberosConnectionParams(
@JsonProperty("server_url") String serverUrl,
@JsonProperty("database") String database,
@JsonProperty("ticket") String ticket) {
super(serverUrl, database);
@JsonProperty("ticket") String ticket,
@JsonProperty("max_connection_lifetime_millis") Long maxConnectionLifetimeMillis,
@JsonProperty("connection_liveness_check_timeout_millis")
Long connectionLivenessCheckTimeoutMillis) {
super(serverUrl, database, maxConnectionLifetimeMillis, connectionLivenessCheckTimeoutMillis);
this.ticket = ticket;
}

Expand Down
8 changes: 6 additions & 2 deletions ...main/java/com/google/cloud/teleport/v2/neo4j/model/connection/NoAuthConnectionParams.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,12 @@ class NoAuthConnectionParams extends ConnectionParams {

@JsonCreator
public NoAuthConnectionParams(
@JsonProperty("server_url") String serverUrl, @JsonProperty("database") String database) {
super(serverUrl, database);
@JsonProperty("server_url") String serverUrl,
@JsonProperty("database") String database,
@JsonProperty("max_connection_lifetime_millis") Long maxConnectionLifetimeMillis,
@JsonProperty("connection_liveness_check_timeout_millis")
Long connectionLivenessCheckTimeoutMillis) {
super(serverUrl, database, maxConnectionLifetimeMillis, connectionLivenessCheckTimeoutMillis);
}

@Override
Expand Down
8 changes: 8 additions & 0 deletions v2/googlecloud-to-neo4j/src/main/resources/schemas/connection.v1.0.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,14 @@
"basic", "none", "kerberos", "bearer", "custom"
],
"default": "basic"
},
"max_connection_lifetime_millis": {
"description": "Max connection lifetime in milliseconds",
"type": "integer"
},
"connection_liveness_check_timeout_millis": {
"description": "Max idle time for pool connections without liveness check upon acquisition",
"type": "integer"
}
},
"required": [
Expand Down
38 changes: 24 additions & 14 deletions ...c/test/java/com/google/cloud/teleport/v2/neo4j/model/connection/ConnectionParamsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,19 +39,21 @@ public void parsesMinimalBasicAuthInformation() throws Exception {

assertThat(connectionParams).isInstanceOf(BasicConnectionParams.class);
assertThat(connectionParams)
.isEqualTo(new BasicConnectionParams("bolt://example.com", null, "neo4j", "password"));
.isEqualTo(
new BasicConnectionParams("bolt://example.com", null, "neo4j", "password", null, null));
}

@Test
public void parsesFullBasicAuthInformation() throws Exception {
ConnectionParams connectionParams =
Json.map(
json(
"{\"auth_type\": \"basic\", \"server_url\": \"bolt://example.com\", \"database\": \"db\", \"username\": \"neo4j\", \"pwd\": \"password\"}"),
"{\"auth_type\": \"basic\", \"server_url\": \"bolt://example.com\", \"database\": \"db\", \"username\": \"neo4j\", \"pwd\": \"password\", \"max_connection_lifetime_millis\": 42, \"connection_liveness_check_timeout_millis\": 4242}"),
ConnectionParams.class);

assertThat(connectionParams)
.isEqualTo(new BasicConnectionParams("bolt://example.com", "db", "neo4j", "password"));
.isEqualTo(
new BasicConnectionParams("bolt://example.com", "db", "neo4j", "password", 42L, 4242L));
}

@Test
Expand All @@ -62,18 +64,20 @@ public void parsesMinimalNoAuthInformation() throws Exception {
ConnectionParams.class);

assertThat(connectionParams).isInstanceOf(NoAuthConnectionParams.class);
assertThat(connectionParams).isEqualTo(new NoAuthConnectionParams("bolt://example.com", null));
assertThat(connectionParams)
.isEqualTo(new NoAuthConnectionParams("bolt://example.com", null, null, null));
}

@Test
public void parsesFullNoAuthInformation() throws Exception {
ConnectionParams connectionParams =
Json.map(
json(
"{\"auth_type\": \"none\", \"server_url\": \"bolt://example.com\", \"database\": \"db\"}"),
"{\"auth_type\": \"none\", \"server_url\": \"bolt://example.com\", \"database\": \"db\", \"max_connection_lifetime_millis\": 42, \"connection_liveness_check_timeout_millis\": 4242}"),
ConnectionParams.class);

assertThat(connectionParams).isEqualTo(new NoAuthConnectionParams("bolt://example.com", "db"));
assertThat(connectionParams)
.isEqualTo(new NoAuthConnectionParams("bolt://example.com", "db", 42L, 4242L));
}

@Test
Expand All @@ -87,20 +91,22 @@ public void parsesMinimalKerberosInformation() throws Exception {
assertThat(connectionParams).isInstanceOf(KerberosConnectionParams.class);
assertThat(connectionParams)
.isEqualTo(
new KerberosConnectionParams("bolt://example.com", null, "dGhpcyBpcyBhIHRpY2tldA=="));
new KerberosConnectionParams(
"bolt://example.com", null, "dGhpcyBpcyBhIHRpY2tldA==", null, null));
}

@Test
public void parsesFullKerberosInformation() throws Exception {
ConnectionParams connectionParams =
Json.map(
json(
"{\"auth_type\": \"kerberos\", \"server_url\": \"bolt://example.com\", \"database\": \"db\", \"ticket\": \"dGhpcyBpcyBhIHRpY2tldA==\"}"),
"{\"auth_type\": \"kerberos\", \"server_url\": \"bolt://example.com\", \"database\": \"db\", \"ticket\": \"dGhpcyBpcyBhIHRpY2tldA==\", \"max_connection_lifetime_millis\": 42, \"connection_liveness_check_timeout_millis\": 4242}"),
ConnectionParams.class);

assertThat(connectionParams)
.isEqualTo(
new KerberosConnectionParams("bolt://example.com", "db", "dGhpcyBpcyBhIHRpY2tldA=="));
new KerberosConnectionParams(
"bolt://example.com", "db", "dGhpcyBpcyBhIHRpY2tldA==", 42L, 4242L));
}

@Test
Expand All @@ -113,19 +119,19 @@ public void parsesMinimalBearerInformation() throws Exception {

assertThat(connectionParams).isInstanceOf(BearerConnectionParams.class);
assertThat(connectionParams)
.isEqualTo(new BearerConnectionParams("bolt://example.com", null, "a-token"));
.isEqualTo(new BearerConnectionParams("bolt://example.com", null, "a-token", null, null));
}

@Test
public void parsesFullBearerInformation() throws Exception {
ConnectionParams connectionParams =
Json.map(
json(
"{\"auth_type\": \"bearer\", \"server_url\": \"bolt://example.com\", \"database\": \"db\", \"token\": \"a-token\"}"),
"{\"auth_type\": \"bearer\", \"server_url\": \"bolt://example.com\", \"database\": \"db\", \"token\": \"a-token\", \"max_connection_lifetime_millis\": 42, \"connection_liveness_check_timeout_millis\": 4242}"),
ConnectionParams.class);

assertThat(connectionParams)
.isEqualTo(new BearerConnectionParams("bolt://example.com", "db", "a-token"));
.isEqualTo(new BearerConnectionParams("bolt://example.com", "db", "a-token", 42L, 4242L));
}

@Test
Expand All @@ -146,6 +152,8 @@ public void parsesMinimalCustomInformation() throws Exception {
"some-credentials",
null,
"a-scheme",
null,
null,
null));
}

Expand All @@ -154,7 +162,7 @@ public void parsesFullCustomInformation() throws Exception {
ConnectionParams connectionParams =
Json.map(
json(
"{\"auth_type\": \"custom\", \"server_url\": \"bolt://example.com\", \"database\": \"db\", \"principal\": \"a-principal\", \"credentials\": \"some-credentials\", \"realm\": \"a-realm\", \"scheme\": \"a-scheme\", \"parameters\": {\"foo\": \"bar\", \"baz\": true}}"),
"{\"auth_type\": \"custom\", \"server_url\": \"bolt://example.com\", \"database\": \"db\", \"principal\": \"a-principal\", \"credentials\": \"some-credentials\", \"realm\": \"a-realm\", \"scheme\": \"a-scheme\", \"parameters\": {\"foo\": \"bar\", \"baz\": true}, \"max_connection_lifetime_millis\": 42, \"connection_liveness_check_timeout_millis\": 4242}"),
ConnectionParams.class);

assertThat(connectionParams)
Expand All @@ -166,7 +174,9 @@ public void parsesFullCustomInformation() throws Exception {
"some-credentials",
"a-realm",
"a-scheme",
Map.of("foo", "bar", "baz", true)));
Map.of("foo", "bar", "baz", true),
42L,
4242L));
}

@NotNull
Expand Down
Loading