Update to latest airbase

Adds proper OpenTracing support to gateway Co-authored-by: Will Morrison <[email protected]>
trinodb · Jan 21, 2025 · c2ba1ea · c2ba1ea
1 parent 1c8ce25
commit c2ba1ea
Show file tree

Hide file tree

Showing 12 changed files with 67 additions and 63 deletions.
diff --git a/gateway-ha/pom.xml b/gateway-ha/pom.xml
@@ -123,11 +123,6 @@
             <artifactId>concurrent</artifactId>
         </dependency>
 
-        <dependency>
-            <groupId>io.airlift</groupId>
-            <artifactId>event</artifactId>
-        </dependency>
-
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>http-client</artifactId>
@@ -183,6 +178,11 @@
             <artifactId>stats</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>tracing</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>units</artifactId>

diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/HaGatewayLauncher.java b/gateway-ha/src/main/java/io/trino/gateway/ha/HaGatewayLauncher.java
@@ -19,7 +19,6 @@
 import com.google.inject.Module;
 import io.airlift.bootstrap.ApplicationConfigurationException;
 import io.airlift.bootstrap.Bootstrap;
-import io.airlift.event.client.EventModule;
 import io.airlift.http.server.HttpServerModule;
 import io.airlift.jaxrs.JaxrsModule;
 import io.airlift.jmx.JmxHttpModule;
@@ -29,6 +28,7 @@
 import io.airlift.log.Logger;
 import io.airlift.node.NodeModule;
 import io.airlift.openmetrics.JmxOpenMetricsModule;
+import io.airlift.tracing.TracingModule;
 import io.airlift.units.Duration;
 import io.trino.gateway.baseapp.BaseApp;
 import io.trino.gateway.ha.config.HaGatewayConfiguration;
@@ -42,6 +42,7 @@
 import static io.trino.gateway.baseapp.BaseApp.addModules;
 import static io.trino.gateway.ha.util.ConfigurationUtils.replaceEnvironmentVariables;
 import static java.lang.String.format;
+import static java.util.Objects.requireNonNullElse;
 
 public class HaGatewayLauncher
 {
@@ -51,10 +52,10 @@ private void start(List<Module> additionalModules, HaGatewayConfiguration config
     {
         long startTime = System.nanoTime();
 
+        String version = requireNonNullElse(HaGatewayLauncher.class.getPackage().getImplementationVersion(), "unknown");
         ImmutableList.Builder<Module> modules = ImmutableList.builder();
         modules.add(
                 new NodeModule(),
-                new EventModule(),
                 new HttpServerModule(),
                 new JmxModule(),
                 new JmxHttpModule(),
@@ -63,6 +64,7 @@ private void start(List<Module> additionalModules, HaGatewayConfiguration config
                 new MBeanModule(),
                 new JsonModule(),
                 new JaxrsModule(),
+                new TracingModule("trino-gateway", version),
                 new BaseApp(configuration));
         modules.addAll(additionalModules);
 

diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/LdapConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/LdapConfiguration.java
@@ -17,8 +17,8 @@
 import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import io.airlift.log.Logger;
 
-import java.io.File;
 import java.io.IOException;
+import java.nio.file.Path;
 
 public class LdapConfiguration
 {
@@ -80,7 +80,7 @@ public static LdapConfiguration load(String path)
     {
         LdapConfiguration configuration = null;
         try {
-            configuration = OBJECT_MAPPER.readValue(new File(path), LdapConfiguration.class);
+            configuration = OBJECT_MAPPER.readValue(Path.of(path).toFile(), LdapConfiguration.class);
         }
         catch (IOException e) {
             log.error(e, "Error loading configuration file");

diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RuleReloadingRoutingGroupSelector.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RuleReloadingRoutingGroupSelector.java
@@ -23,7 +23,6 @@
 import org.jeasy.rules.mvel.MVELRuleFactory;
 import org.jeasy.rules.support.reader.YamlRuleDefinitionReader;
 
-import java.io.FileReader;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.attribute.BasicFileAttributes;
@@ -40,7 +39,7 @@ public class RuleReloadingRoutingGroupSelector
     private static final Logger log = Logger.get(RuleReloadingRoutingGroupSelector.class);
     private final RulesEngine rulesEngine = new DefaultRulesEngine();
     private final MVELRuleFactory ruleFactory = new MVELRuleFactory(new YamlRuleDefinitionReader());
-    private final String rulesConfigPath;
+    private final Path rulesConfigPath;
     private volatile Rules rules = new Rules();
     private volatile long lastUpdatedTime;
     private final ReadWriteLock readWriteLock = new ReentrantReadWriteLock(true);
@@ -49,13 +48,13 @@ public class RuleReloadingRoutingGroupSelector
 
     RuleReloadingRoutingGroupSelector(String rulesConfigPath, RequestAnalyzerConfig requestAnalyzerConfig)
     {
-        this.rulesConfigPath = rulesConfigPath;
+        this.rulesConfigPath = Path.of(rulesConfigPath);
         this.requestAnalyzerConfig = requestAnalyzerConfig;
         trinoRequestUserProvider = new TrinoRequestUser.TrinoRequestUserProvider(requestAnalyzerConfig);
         try {
             rules = ruleFactory.createRules(
-                    new FileReader(rulesConfigPath, UTF_8));
-            BasicFileAttributes attr = Files.readAttributes(Path.of(rulesConfigPath),
+                    Files.newBufferedReader(this.rulesConfigPath, UTF_8));
+            BasicFileAttributes attr = Files.readAttributes(this.rulesConfigPath,
                     BasicFileAttributes.class);
             lastUpdatedTime = attr.lastModifiedTime().toMillis();
         }
@@ -70,7 +69,7 @@ public class RuleReloadingRoutingGroupSelector
     public String findRoutingGroup(HttpServletRequest request)
     {
         try {
-            BasicFileAttributes attr = Files.readAttributes(Path.of(rulesConfigPath),
+            BasicFileAttributes attr = Files.readAttributes(rulesConfigPath,
                     BasicFileAttributes.class);
             log.debug("File modified time: %s. lastUpdatedTime: %s", attr.lastModifiedTime(), lastUpdatedTime);
             if (attr.lastModifiedTime().toMillis() > lastUpdatedTime) {
@@ -82,7 +81,7 @@ public String findRoutingGroup(HttpServletRequest request)
                         // thread finds the condition true and acquires the lock before this one
                         log.info("Updating rules to file modified at %s", attr.lastModifiedTime());
                         rules = ruleFactory.createRules(
-                                new FileReader(rulesConfigPath, UTF_8));
+                                Files.newBufferedReader(rulesConfigPath, UTF_8));
                         lastUpdatedTime = attr.lastModifiedTime().toMillis();
                     }
                 }

diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/security/LbKeyProvider.java b/gateway-ha/src/main/java/io/trino/gateway/ha/security/LbKeyProvider.java
@@ -17,7 +17,9 @@
 import org.bouncycastle.util.io.pem.PemObject;
 import org.bouncycastle.util.io.pem.PemReader;
 
-import java.io.FileReader;
+import java.io.BufferedReader;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.security.KeyFactory;
 import java.security.NoSuchAlgorithmException;
 import java.security.PrivateKey;
@@ -49,7 +51,7 @@ public LbKeyProvider(SelfSignKeyPairConfiguration keypairConfig)
 
         try {
             String publicKeyRsa = keypairConfig.publicKeyRsa();
-            try (FileReader keyReader = new FileReader(publicKeyRsa, UTF_8);
+            try (BufferedReader keyReader = Files.newBufferedReader(Path.of(publicKeyRsa), UTF_8);
                     PemReader pemReader = new PemReader(keyReader)) {
                 PemObject pemObject = pemReader.readPemObject();
                 byte[] content = pemObject.getContent();
@@ -58,7 +60,7 @@ public LbKeyProvider(SelfSignKeyPairConfiguration keypairConfig)
             }
 
             String privateKeyRsa = keypairConfig.privateKeyRsa();
-            try (FileReader keyReader = new FileReader(privateKeyRsa, UTF_8);
+            try (BufferedReader keyReader = Files.newBufferedReader(Path.of(privateKeyRsa), UTF_8);
                     PemReader pemReader = new PemReader(keyReader)) {
                 PemObject pemObject = pemReader.readPemObject();
                 byte[] content = pemObject.getContent();

diff --git a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java
@@ -179,7 +179,6 @@ private void performRequest(
         cookieBuilder.addAll(getOAuth2GatewayCookie(remoteUri, servletRequest));
 
         Request request = requestBuilder
-                .setPreserveAuthorizationOnRedirect(true)
                 .setFollowRedirects(false)
                 .build();
 

diff --git a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyResponseHandler.java b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyResponseHandler.java
@@ -22,7 +22,6 @@
 
 import java.io.IOException;
 
-import static com.google.common.io.ByteStreams.toByteArray;
 import static java.util.Objects.requireNonNull;
 
 public class ProxyResponseHandler
@@ -38,7 +37,7 @@ public ProxyResponse handleException(Request request, Exception exception)
     public ProxyResponse handle(Request request, Response response)
     {
         try {
-            return new ProxyResponse(response.getStatusCode(), response.getHeaders(), toByteArray(response.getInputStream()));
+            return new ProxyResponse(response.getStatusCode(), response.getHeaders(), response.getInputStream().readAllBytes());
         }
         catch (IOException e) {
             throw new ProxyException("Failed reading response from remote Trino server", e);

diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java b/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java
@@ -28,12 +28,13 @@
 import org.jdbi.v3.core.Handle;
 import org.jdbi.v3.core.Jdbi;
 
+import java.io.BufferedWriter;
 import java.io.File;
-import java.io.FileWriter;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URL;
-import java.nio.file.Paths;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.Random;
 import java.util.Scanner;
 import java.util.concurrent.TimeUnit;
@@ -66,21 +67,20 @@ public static void seedRequiredData(TestConfig testConfig)
     }
 
     public static void prepareMockBackend(
-            MockWebServer backend, int customBackendPort, String expectedResonse)
+            MockWebServer backend, int customBackendPort, String expectedResponse)
             throws IOException
     {
         backend.start(customBackendPort);
         backend.enqueue(new MockResponse()
-                .setBody(expectedResonse)
+                .setBody(expectedResponse)
                 .addHeader(CONTENT_ENCODING, PLAIN_TEXT_UTF_8)
                 .setResponseCode(200));
     }
 
     public static TestConfig buildGatewayConfigAndSeedDb(int routerPort, String configFile)
             throws Exception
     {
-        File baseDir = new File(System.getProperty("java.io.tmpdir"));
-        File tempH2DbDir = new File(baseDir, "h2db-" + RANDOM.nextInt() + System.currentTimeMillis());
+        File tempH2DbDir = Path.of(System.getProperty("java.io.tmpdir"), "h2db-" + RANDOM.nextInt() + System.currentTimeMillis()).toFile();
         tempH2DbDir.deleteOnExit();
 
         URL resource = HaGatewayTestUtils.class.getClassLoader().getResource("auth/localhost.jks");
@@ -91,14 +91,15 @@ public static TestConfig buildGatewayConfigAndSeedDb(int routerPort, String conf
                         .replace(
                                 "APPLICATION_CONNECTOR_PORT", String.valueOf(30000 + (int) (Math.random() * 1000)))
                         .replace("ADMIN_CONNECTOR_PORT", String.valueOf(31000 + (int) (Math.random() * 1000)))
-                        .replace("LOCALHOST_JKS", Paths.get(resource.toURI()).toFile().getAbsolutePath())
-                        .replace("RESOURCES_DIR", Paths.get("src", "test", "resources").toFile().getAbsolutePath());
+                        .replace("LOCALHOST_JKS", Path.of(resource.toURI()).toString())
+                        .replace("RESOURCES_DIR", Path.of("src", "test", "resources").toAbsolutePath().toString());
 
         File target = File.createTempFile("config-" + System.currentTimeMillis(), "config.yaml");
 
-        FileWriter fw = new FileWriter(target, UTF_8);
-        fw.append(configStr);
-        fw.flush();
+        try (BufferedWriter writer = Files.newBufferedWriter(target.toPath(), UTF_8)) {
+            writer.append(configStr);
+        }
+
         log.info("Test Gateway Config \n[%s]", configStr);
         TestConfig testConfig = new TestConfig(target.getAbsolutePath(), tempH2DbDir.getAbsolutePath());
         seedRequiredData(testConfig);

diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestingJdbcConnectionManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestingJdbcConnectionManager.java
@@ -18,15 +18,15 @@
 import org.jdbi.v3.core.Jdbi;
 
 import java.io.File;
+import java.nio.file.Path;
 
 public final class TestingJdbcConnectionManager
 {
     private TestingJdbcConnectionManager() {}
 
     public static JdbcConnectionManager createTestingJdbcConnectionManager()
     {
-        File baseDir = new File(System.getProperty("java.io.tmpdir"));
-        File tempH2DbDir = new File(baseDir, "h2db-" + System.currentTimeMillis());
+        File tempH2DbDir = Path.of(System.getProperty("java.io.tmpdir"), "h2db-" + System.currentTimeMillis()).toFile();
         tempH2DbDir.deleteOnExit();
         String jdbcUrl = "jdbc:h2:" + tempH2DbDir.getAbsolutePath();
         HaGatewayTestUtils.seedRequiredData(new HaGatewayTestUtils.TestConfig("", tempH2DbDir.getAbsolutePath()));

diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java
@@ -27,12 +27,13 @@
 import org.junit.jupiter.params.provider.MethodSource;
 
 import java.io.BufferedReader;
+import java.io.BufferedWriter;
 import java.io.File;
-import java.io.FileReader;
-import java.io.FileWriter;
 import java.io.IOException;
 import java.io.Reader;
 import java.io.StringReader;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.Base64;
 import java.util.Collections;
@@ -251,16 +252,16 @@ void testByRoutingRulesEngineFileChange()
     {
         File file = File.createTempFile("routing_rules", ".yml");
 
-        FileWriter fw = new FileWriter(file, UTF_8);
-        fw.write(
-                "---\n"
-                        + "name: \"airflow1\"\n"
-                        + "description: \"original rule\"\n"
-                        + "condition: \"request.getHeader(\\\"X-Trino-Source\\\") == \\\"airflow\\\"\"\n"
-                        + "actions:\n"
-                        + "  - \"result.put(\\\"routingGroup\\\", \\\"etl\\\")\"");
-        fw.close();
-        long lastModifed = file.lastModified();
+        try (BufferedWriter writer = Files.newBufferedWriter(file.toPath(), UTF_8)) {
+            writer.write(
+                    "---\n"
+                            + "name: \"airflow1\"\n"
+                            + "description: \"original rule\"\n"
+                            + "condition: \"request.getHeader(\\\"X-Trino-Source\\\") == \\\"airflow\\\"\"\n"
+                            + "actions:\n"
+                            + "  - \"result.put(\\\"routingGroup\\\", \\\"etl\\\")\"");
+        }
+        long lastModified = file.lastModified();
 
         RoutingGroupSelector routingGroupSelector =
                 RoutingGroupSelector.byRoutingRulesEngine(file.getPath(), requestAnalyzerConfig);
@@ -271,16 +272,16 @@ void testByRoutingRulesEngineFileChange()
         assertThat(routingGroupSelector.findRoutingGroup(mockRequest))
                 .isEqualTo("etl");
 
-        fw = new FileWriter(file, UTF_8);
-        fw.write(
-                "---\n"
-                        + "name: \"airflow2\"\n"
-                        + "description: \"updated rule\"\n"
-                        + "condition: \"request.getHeader(\\\"X-Trino-Source\\\") == \\\"airflow\\\"\"\n"
-                        + "actions:\n"
-                        + "  - \"result.put(\\\"routingGroup\\\", \\\"etl2\\\")\""); // change from etl to etl2
-        fw.close();
-        assertThat(file.setLastModified(lastModifed + 1000)).isTrue();
+        try (BufferedWriter writer = Files.newBufferedWriter(file.toPath(), UTF_8)) {
+            writer.write(
+                    "---\n"
+                            + "name: \"airflow2\"\n"
+                            + "description: \"updated rule\"\n"
+                            + "condition: \"request.getHeader(\\\"X-Trino-Source\\\") == \\\"airflow\\\"\"\n"
+                            + "actions:\n"
+                            + "  - \"result.put(\\\"routingGroup\\\", \\\"etl2\\\")\""); // change from etl to etl2
+        }
+        assertThat(file.setLastModified(lastModified + 1000)).isTrue();
 
         when(mockRequest.getHeader(TRINO_SOURCE_HEADER)).thenReturn("airflow");
         assertThat(routingGroupSelector.findRoutingGroup(mockRequest))
@@ -450,7 +451,7 @@ private HttpServletRequest prepareMockRequest()
     void testLongQuery()
             throws IOException
     {
-        BufferedReader bufferedReader = new BufferedReader(new FileReader("src/test/resources/wide_select.sql", UTF_8));
+        BufferedReader bufferedReader = Files.newBufferedReader(Path.of("src/test/resources/wide_select.sql"), UTF_8);
         HttpServletRequest mockRequest = prepareMockRequest();
         when(mockRequest.getReader()).thenReturn(bufferedReader);
         TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(

diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestSpecificDbResourceGroupsManager.java
@@ -23,6 +23,7 @@
 import org.junit.jupiter.api.TestInstance.Lifecycle;
 
 import java.io.File;
+import java.nio.file.Path;
 import java.util.List;
 
 import static io.trino.gateway.ha.router.ResourceGroupsManager.ResourceGroupsDetail;
@@ -40,8 +41,7 @@ final class TestSpecificDbResourceGroupsManager
     void setUp()
     {
         specificDb = "h2db-" + System.currentTimeMillis();
-        File baseDir = new File(System.getProperty("java.io.tmpdir"));
-        File tempH2DbDir = new File(baseDir, specificDb);
+        File tempH2DbDir = Path.of(System.getProperty("java.io.tmpdir"), specificDb).toFile();
         tempH2DbDir.deleteOnExit();
         String jdbcUrl = "jdbc:h2:" + tempH2DbDir.getAbsolutePath();
         HaGatewayTestUtils.seedRequiredData(