From 2182c4b43253cfd0de0df4b2e1f0d875d55afbc9 Mon Sep 17 00:00:00 2001
From: Tadaya Tsuyukubo <tadaya@ttddyy.net>
Date: Tue, 1 Aug 2023 18:55:22 -0700
Subject: [PATCH] Allow null host param on ObservationProxyExecutionListener
 constructor

Relates #135
---
 .../ObservationProxyExecutionListener.java    |  5 +++-
 ...ObservationProxyExecutionListenerTest.java | 26 ++++++++++++++-----
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/main/java/io/r2dbc/proxy/observation/ObservationProxyExecutionListener.java b/src/main/java/io/r2dbc/proxy/observation/ObservationProxyExecutionListener.java
index 963cb8cd..fa03d425 100644
--- a/src/main/java/io/r2dbc/proxy/observation/ObservationProxyExecutionListener.java
+++ b/src/main/java/io/r2dbc/proxy/observation/ObservationProxyExecutionListener.java
@@ -94,7 +94,7 @@ public ObservationProxyExecutionListener(ObservationRegistry observationRegistry
      * @since 1.1.2
      */
     public ObservationProxyExecutionListener(ObservationRegistry observationRegistry,
-                                             ConnectionFactory connectionFactory, String host, @Nullable Integer port) {
+                                             ConnectionFactory connectionFactory, @Nullable String host, @Nullable Integer port) {
         this.observationRegistry = observationRegistry;
         this.connectionFactory = connectionFactory;
         this.remoteServiceAddress = buildRemoteServiceAddress(host, port);
@@ -110,6 +110,9 @@ private String parseR2dbcConnectionUrl(String connectionUrl) {
 
     @Nullable
     private String buildRemoteServiceAddress(@Nullable String host, @Nullable Integer portNumber) {
+        if (host == null) {
+            return null;
+        }
         int port = portNumber != null ? portNumber : -1;
         try {
             URI uri = new URI(null, null, host, port, null, null, null);
diff --git a/src/test/java/io/r2dbc/proxy/observation/ObservationProxyExecutionListenerTest.java b/src/test/java/io/r2dbc/proxy/observation/ObservationProxyExecutionListenerTest.java
index 9f90ac34..b2aa8556 100644
--- a/src/test/java/io/r2dbc/proxy/observation/ObservationProxyExecutionListenerTest.java
+++ b/src/test/java/io/r2dbc/proxy/observation/ObservationProxyExecutionListenerTest.java
@@ -28,6 +28,7 @@
 import io.micrometer.tracing.propagation.Propagator;
 import io.micrometer.tracing.test.simple.SimpleSpan;
 import io.micrometer.tracing.test.simple.SimpleTracer;
+import io.micrometer.tracing.test.simple.SpanAssert;
 import io.micrometer.tracing.test.simple.TracingAssertions;
 import io.r2dbc.proxy.core.Bindings;
 import io.r2dbc.proxy.core.BoundValue;
@@ -41,6 +42,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import reactor.util.annotation.Nullable;
 import reactor.util.context.Context;
 import reactor.util.context.ContextView;
 
@@ -195,7 +197,7 @@ private static Stream<Arguments> remoteServiceAddressWithConnectionUrl() {
 
     @ParameterizedTest
     @MethodSource
-    void remoteServiceAddressWithConnectionHostAndPort(String host, Integer port, String expectedIp, int expectedPort) {
+    void remoteServiceAddressWithConnectionHostAndPort(String host, Integer port, String expectedIp, Integer expectedPort) {
         runAndVerifyRemoteServiceAddress(() -> {
             ConnectionFactory connectionFactory = createMockConnectionFactory();
             return new ObservationProxyExecutionListener(this.registry, connectionFactory, host, port);
@@ -206,11 +208,14 @@ private static Stream<Arguments> remoteServiceAddressWithConnectionHostAndPort()
         return Stream.of(
             Arguments.of("192.168.1.1", 5432, "192.168.1.1", 5432),
             Arguments.of("192.168.1.1", null, "192.168.1.1", -1),
-            Arguments.of("192.168.1.1", -1, "192.168.1.1", -1)
+            Arguments.of("192.168.1.1", -1, "192.168.1.1", -1),
+            Arguments.of(null, 5432, null, null),
+            Arguments.of(null, -1, null, null),
+            Arguments.of(null, null, null, null)
         );
     }
 
-    void runAndVerifyRemoteServiceAddress(Supplier<ObservationProxyExecutionListener> listenerSupplier, String expectedIp, int expectedPort) {
+    void runAndVerifyRemoteServiceAddress(Supplier<ObservationProxyExecutionListener> listenerSupplier, @Nullable String expectedIp, @Nullable Integer expectedPort) {
         this.registry.observationConfig().observationHandler(new PropagatingSenderTracingObservationHandler<QueryContext>(this.tracer, NOOP_PROPAGATOR));
         QueryExecutionInfo queryExecutionInfo = createQueryExecutionInfo();
 
@@ -218,9 +223,18 @@ void runAndVerifyRemoteServiceAddress(Supplier<ObservationProxyExecutionListener
         listener.beforeQuery(queryExecutionInfo);
         listener.afterQuery(queryExecutionInfo);
 
-        TracingAssertions.assertThat(this.tracer).onlySpan()
-            .hasIpEqualTo(expectedIp)
-            .hasPortEqualTo(expectedPort);
+        TracingAssertions.assertThat(this.tracer).onlySpan();
+        SimpleSpan span = this.tracer.onlySpan();
+        if (expectedIp == null) {
+            SpanAssert.assertThat(span).hasIpThatIsBlank();
+        } else {
+            SpanAssert.assertThat(span).hasIpEqualTo(expectedIp);
+        }
+        if (expectedPort == null) {
+            SpanAssert.assertThat(span).hasPortThatIsNotSet();
+        } else {
+            SpanAssert.assertThat(span).hasPortEqualTo(expectedPort);
+        }
     }
 
 }