Skip to content

Commit b5b75d2

Browse files
sergiyvamzsergiyvamz
committed
fix: connection identification and tracking
1 parent 1140de8 commit b5b75d2

File tree

6 files changed

+54
-14
lines changed

6 files changed

+54
-14
lines changed

‎wrapper/src/main/java/software/amazon/jdbc/HostSpec.java

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import java.util.concurrent.ConcurrentHashMap;
2626
import software.amazon.jdbc.hostavailability.HostAvailability;
2727
import software.amazon.jdbc.hostavailability.HostAvailabilityStrategy;
28+
import software.amazon.jdbc.util.StringUtils;
2829

2930
/**
3031
* An object representing connection info for a given host. Modifiable fields are thread-safe to support sharing this
@@ -150,8 +151,10 @@ public void addAlias(final String... alias) {
150151
}
151152

152153
Arrays.asList(alias).forEach(x -> {
153-
this.aliases.add(x);
154-
this.allAliases.add(x);
154+
if (!StringUtils.isNullOrEmpty(x)) {
155+
this.aliases.add(x);
156+
this.allAliases.add(x);
157+
}
155158
});
156159
}
157160

@@ -160,8 +163,10 @@ public void removeAlias(final String... alias) {
160163
return;
161164
}
162165
Arrays.asList(alias).forEach(x -> {
163-
this.aliases.remove(x);
164-
this.allAliases.remove(x);
166+
if (!StringUtils.isNullOrEmpty(x)) {
167+
this.aliases.remove(x);
168+
this.allAliases.remove(x);
169+
}
165170
});
166171
}
167172

‎wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -633,6 +633,8 @@ public void fillAliases(Connection connection, HostSpec hostSpec) throws SQLExce
633633
final HostSpec host = this.identifyConnection(connection);
634634
if (host != null) {
635635
hostSpec.addAlias(host.asAliases().toArray(new String[]{}));
636+
} else {
637+
LOGGER.finest("Can't identify connection");
636638
}
637639
}
638640

‎wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/RdsHostListProvider.java

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -648,17 +648,43 @@ public HostSpec identifyConnection(Connection connection) throws SQLException {
648648
if (resultSet.next()) {
649649
final String instanceName = resultSet.getString(1);
650650

651-
final List<HostSpec> topology = this.refresh();
651+
List<HostSpec> topology = this.refresh(connection);
652652

653+
boolean isForcedRefresh = false;
653654
if (topology == null) {
655+
LOGGER.finest("forceRefresh topology");
656+
topology = this.forceRefresh(connection);
657+
isForcedRefresh = true;
658+
}
659+
660+
if (topology == null) {
661+
LOGGER.finest("Topology is null");
654662
return null;
655663
}
656664

657-
return topology
665+
HostSpec foundHost = topology
658666
.stream()
659667
.filter(host -> Objects.equals(instanceName, host.getHostId()))
660668
.findAny()
661669
.orElse(null);
670+
671+
if (foundHost == null && !isForcedRefresh) {
672+
LOGGER.finest("forceRefresh topology 2");
673+
topology = this.forceRefresh(connection);
674+
675+
if (topology == null) {
676+
LOGGER.finest("Topology is null 2");
677+
return null;
678+
}
679+
680+
foundHost = topology
681+
.stream()
682+
.filter(host -> Objects.equals(instanceName, host.getHostId()))
683+
.findAny()
684+
.orElse(null);
685+
}
686+
687+
return foundHost;
662688
}
663689
} catch (final SQLException e) {
664690
throw new SQLException(Messages.get("RdsHostListProvider.errorIdentifyConnection"), e);

‎wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ public boolean isRdsDns(final String host) {
141141
}
142142

143143
public boolean isRdsInstance(final String host) {
144-
return getDnsGroup(host) == null;
144+
return getDnsGroup(host) == null && isRdsDns(host);
145145
}
146146

147147
public boolean isRdsProxyDns(final String host) {

‎wrapper/src/test/java/software/amazon/jdbc/hostlistprovider/RdsHostListProviderTest.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -473,7 +473,8 @@ void testIdentifyConnectionNullTopology() throws SQLException {
473473

474474
when(mockResultSet.next()).thenReturn(true);
475475
when(mockResultSet.getString(eq(1))).thenReturn("instance-1");
476-
when(rdsHostListProvider.refresh(eq(mockConnection))).thenReturn(null);
476+
doReturn(null).when(rdsHostListProvider).refresh(mockConnection);
477+
doReturn(null).when(rdsHostListProvider).forceRefresh(mockConnection);
477478

478479
assertNull(rdsHostListProvider.identifyConnection(mockConnection));
479480
}
@@ -492,7 +493,8 @@ void testIdentifyConnectionHostNotInTopology() throws SQLException {
492493
"jdbc:someprotocol://url"));
493494
when(mockResultSet.next()).thenReturn(true);
494495
when(mockResultSet.getString(eq(1))).thenReturn("instance-1");
495-
when(rdsHostListProvider.refresh(eq(mockConnection))).thenReturn(cachedTopology);
496+
doReturn(cachedTopology).when(rdsHostListProvider).refresh(mockConnection);
497+
doReturn(cachedTopology).when(rdsHostListProvider).forceRefresh(mockConnection);
496498

497499
assertNull(rdsHostListProvider.identifyConnection(mockConnection));
498500
}
@@ -512,7 +514,8 @@ void testIdentifyConnectionHostInTopology() throws SQLException {
512514
"jdbc:someprotocol://url"));
513515
when(mockResultSet.next()).thenReturn(true);
514516
when(mockResultSet.getString(eq(1))).thenReturn("instance-a-1");
515-
when(rdsHostListProvider.refresh()).thenReturn(cachedTopology);
517+
doReturn(cachedTopology).when(rdsHostListProvider).refresh(mockConnection);
518+
doReturn(cachedTopology).when(rdsHostListProvider).forceRefresh(mockConnection);
516519

517520
final HostSpec actual = rdsHostListProvider.identifyConnection(mockConnection);
518521
assertEquals("instance-a-1.xyz.us-east-2.rds.amazonaws.com", actual.getHost());

‎wrapper/src/test/java/software/amazon/jdbc/hostlistprovider/RdsMultiAzDbClusterListProviderTest.java

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import java.util.concurrent.TimeUnit;
4747
import org.junit.jupiter.api.AfterEach;
4848
import org.junit.jupiter.api.BeforeEach;
49+
import org.junit.jupiter.api.Disabled;
4950
import org.junit.jupiter.api.Test;
5051
import org.mockito.ArgumentCaptor;
5152
import org.mockito.Captor;
@@ -454,7 +455,8 @@ void testIdentifyConnectionNullTopology() throws SQLException {
454455

455456
when(mockResultSet.next()).thenReturn(true);
456457
when(mockResultSet.getString(eq(1))).thenReturn("instance-1");
457-
when(rdsMazDbClusterHostListProvider.refresh(eq(mockConnection))).thenReturn(null);
458+
doReturn(null).when(rdsMazDbClusterHostListProvider).refresh(mockConnection);
459+
doReturn(null).when(rdsMazDbClusterHostListProvider).forceRefresh(mockConnection);
458460

459461
assertNull(rdsMazDbClusterHostListProvider.identifyConnection(mockConnection));
460462
}
@@ -473,7 +475,8 @@ void testIdentifyConnectionHostNotInTopology() throws SQLException {
473475
"jdbc:someprotocol://url"));
474476
when(mockResultSet.next()).thenReturn(true);
475477
when(mockResultSet.getString(eq(1))).thenReturn("instance-1");
476-
when(rdsMazDbClusterHostListProvider.refresh(eq(mockConnection))).thenReturn(cachedTopology);
478+
doReturn(cachedTopology).when(rdsMazDbClusterHostListProvider).refresh(mockConnection);
479+
doReturn(cachedTopology).when(rdsMazDbClusterHostListProvider).forceRefresh(mockConnection);
477480

478481
assertNull(rdsMazDbClusterHostListProvider.identifyConnection(mockConnection));
479482
}
@@ -482,18 +485,19 @@ void testIdentifyConnectionHostNotInTopology() throws SQLException {
482485
void testIdentifyConnectionHostInTopology() throws SQLException {
483486
final HostSpec expectedHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
484487
.host("instance-a-1.xyz.us-east-2.rds.amazonaws.com")
488+
.hostId("instance-a-1")
485489
.port(HostSpec.NO_PORT)
486490
.role(HostRole.WRITER)
487491
.build();
488-
expectedHost.setHostId("instance-a-1");
489492
final List<HostSpec> cachedTopology = Collections.singletonList(expectedHost);
490493

491494
rdsMazDbClusterHostListProvider = Mockito.spy(getRdsMazDbClusterHostListProvider(
492495
mockHostListProviderService,
493496
"jdbc:someprotocol://url"));
494497
when(mockResultSet.next()).thenReturn(true);
495498
when(mockResultSet.getString(eq(1))).thenReturn("instance-a-1");
496-
when(rdsMazDbClusterHostListProvider.refresh()).thenReturn(cachedTopology);
499+
doReturn(cachedTopology).when(rdsMazDbClusterHostListProvider).refresh(mockConnection);
500+
doReturn(cachedTopology).when(rdsMazDbClusterHostListProvider).forceRefresh(mockConnection);
497501

498502
final HostSpec actual = rdsMazDbClusterHostListProvider.identifyConnection(mockConnection);
499503
assertEquals("instance-a-1.xyz.us-east-2.rds.amazonaws.com", actual.getHost());

0 commit comments

Comments
 (0)