Add support for mTLS authentication in Arrow Flight client by elbinpallimalilibm · Pull Request #25179 · prestodb/presto

elbinpallimalilibm · 2025-05-23T02:50:09Z

Description

Add support for mTLS authentication in Arrow Flight client

Motivation and Context

If the Flight server has mTLS authentication enabled, then the Flight client should be able to use client certificate and key.

Impact

Test Plan

Added positive and negative test cases against an mTLS enabled Flight server.

Contributor checklist

Please make sure your submission complies with our contributing guide, in particular code style and commit standards.
PR description addresses the issue accurately and concisely. If the change is non-trivial, a GitHub Issue is referenced.
Documented new properties (with its default value), SQL syntax, functions, or other functionality.
If release notes are required, they follow the release notes guidelines.
Adequate tests were added if applicable.
CI passed.

Release Notes

Please follow release notes guidelines and fill in the release notes below.

== RELEASE NOTES ==

Arrow Flight Connector Template Changes
* Added support for mTLS authentication in Arrow Flight client.

elbinpallimalilibm · 2025-05-23T03:05:02Z

https://github.com/prestodb/presto/actions/runs/15201290385/job/42755765576?pr=25179 @steveburnett I'm not able to figure out why the presto-docs check is failing. Can you help here?

Edit : Waiting for fix to be merged here.

steveburnett · 2025-05-23T18:56:27Z

Fix has been merged, and the test is now passing - see #25188 for an example. Rebase your PR to re-run the CI tests and docs / test (:presto-docs) (pull_request) should pass for this PR as well.

pratyakshsharma

Thank you for the fix @elbinpallimalilibm. I have added few comments, please check.

pratyakshsharma · 2025-05-27T08:30:56Z

presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowFlightConfig.java

+        return flightClientSSLCertificate;
+    }
+
+    @Config("arrow-flight.client-ssl-certificate")


I think it will be good to add a comment here saying this is needed for mTLS auth.

Added comments.

instead of a javadoc, could you put this in a @ConfigDescription?

pratyakshsharma · 2025-05-27T08:32:30Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+            Optional<InputStream> clientCertificate = Optional.empty();
+            Optional<InputStream> clientKey = Optional.empty();
+            if (config.getFlightClientSSLCertificate() != null && config.getFlightClientSSLKey() != null) {
+                clientCertificate = Optional.of(newInputStream(Paths.get(config.getFlightClientSSLCertificate())));


trying to understand it a bit better, what happens if the certificate and key are invalid? Lets use a try-catch block here? And maybe add a test case covering this scenario?

If the cert is invalid, executing a query will give the user an error that the cert is invalid. Added a test case that covers this scenario.

Lets add a try-catch here as well and modify the error message in the test case accordingly. Thank you for adding test case for this though.

The invalid cert exception is thrown only at line 84 FlightClient flightClient = flightClientBuilder.build(); and we might get exception due to other reasons as well from the build method. So adding a try...catch here will not help in modifying the error message.

Ok, I actually foresee improperly configured client cert/key as a very probable source of error, and hence wanted to cover the scenario with a proper user facing message. Anyways I leave the final decision to you.

Refactored the code to catch errors due to invalid cert or key file. Rethrowing a Presto Exception with a custom message for those scenarios.

pratyakshsharma · 2025-05-27T08:38:15Z

presto-docs/src/main/sphinx/connector/base-arrow-flight.rst

 ``arrow-flight.server.port``                Flight server port
-``arrow-flight.server-ssl-certificate``     Pass ssl certificate
+``arrow-flight.server-ssl-certificate``     Path to SSL certificate of Flight server
+``arrow-flight.client-ssl-certificate``     Path to SSL certificate that Flight client should use


nit: lets modify these to include "in case of mTLS authentication" to make it more clear?

Updated the docs.

pratyakshsharma · 2025-05-27T08:39:42Z

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMTLSFails.java

+    }
+
+    @BeforeClass
+    public void setup()


The 2 test classes have a lot of duplicate code. Please see if we can use inheritance to avoid this redundant code?

Refactored the test classes.

steveburnett

LGTM! (docs)

Pull branch, local doc build, looks good. Thanks!

pratyakshsharma

Thank you for addressing earlier comments. I have few more minor comments, once addressed, I will approve.

pratyakshsharma · 2025-05-28T09:02:38Z

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java

    }

-    private static DistributedQueryRunner createQueryRunner(
+    public static DistributedQueryRunner createQueryRunner(


why are we modifying the access-modifier?

We need to create query runner with pre-defined catalog properties from the new test classes. That's why the method was changed to public from private.

Does protected or other modifier work instead?

Changed to protected

pratyakshsharma · 2025-05-28T09:05:14Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+            Optional<InputStream> clientCertificate = Optional.empty();
+            Optional<InputStream> clientKey = Optional.empty();
+            if (config.getFlightClientSSLCertificate() != null && config.getFlightClientSSLKey() != null) {
+                clientCertificate = Optional.of(newInputStream(Paths.get(config.getFlightClientSSLCertificate())));


Lets add a try-catch here as well and modify the error message in the test case accordingly. Thank you for adding test case for this though.

pratyakshsharma · 2025-05-28T09:12:52Z

...row-flight/src/test/java/com/facebook/plugin/arrow/AbstractArrowFlightMTLSTestFramework.java

+    }
+
+    @BeforeClass
+    public void setup()


Does this method need to be public? How about making it package private? Similarly for other methods apart from createQueryRunner, lets make them as restricted as possible.

pratyakshsharma · 2025-05-28T09:13:29Z

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMTLS.java

+    }
+
+    @Override
+    protected Map<String, String> getCatalogProperties()


The access modifiers here will change based on my above comment.

pratyakshsharma

Couple more comments, rest looks good

pratyakshsharma · 2025-05-28T12:50:04Z

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java

    }

-    private static DistributedQueryRunner createQueryRunner(
+    public static DistributedQueryRunner createQueryRunner(


Does protected or other modifier work instead?

pratyakshsharma · 2025-05-28T12:50:24Z

...row-flight/src/test/java/com/facebook/plugin/arrow/AbstractArrowFlightMTLSTestFramework.java

+
+    abstract Map<String, String> getCatalogProperties();
+
+    protected int getServerPort()


Lets change this to package private as well?

pratyakshsharma

Thank you for patiently addressing all comments. Few refactorings needed, rest looks good

pratyakshsharma · 2025-05-28T13:30:25Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

        catch (Exception e) {
-            throw new ArrowException(ARROW_FLIGHT_CLIENT_ERROR, "Error creating flight client: " + e.getMessage(), e);
+            Optional<Throwable> cause = Optional.ofNullable(e.getCause());
+            if (cause.filter(c -> c instanceof InvalidKeyException).isPresent()) {


Suggested change

if (cause.filter(c -> c instanceof InvalidKeyException).isPresent()) {

if (e instanceOf InvalidKeyException) {

can be simplified like this.

e will be instance of IllegalArgumentException. Inner exception e.getCause if not null, will be an instance of InvalidKeyException

pratyakshsharma · 2025-05-28T13:31:07Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+            if (cause.filter(c -> c instanceof InvalidKeyException).isPresent()) {
+                throw new ArrowException(ARROW_FLIGHT_INVALID_KEY_ERROR, "Error creating flight client, invalid key file: " + e.getMessage(), e);
+            }
+            else if (cause.filter(c -> c instanceof CertificateException).isPresent()) {


pratyakshsharma · 2025-05-28T13:32:26Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+                    clientCertificate.get().close();
+                }
+                catch (IOException e) {
+                    logger.error("Error closing input stream", e);


the error messages here are all identical, can be modified in different blocks.

pratyakshsharma · 2025-05-28T13:32:49Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+                    clientCertificate.get().close();
+                }
+                catch (IOException e) {
+                    logger.error("Error closing input stream", e);


Suggested change

logger.error("Error closing input stream", e);

logger.error("Error closing input stream for clientCertificate", e);

pratyakshsharma · 2025-05-28T13:33:00Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

+                    clientKey.get().close();
+                }
+                catch (IOException e) {
+                    logger.error("Error closing input stream", e);


pratyakshsharma

Thank you for the proactive responses. LGTM!

pratyakshsharma · 2025-05-28T13:47:18Z

@prestodb/committers this should be good for final pass

ZacBlanco

One comment. Also, I think instead of having 3 test classes, let's combine the TestArrowFlightMTLSFails and TestArrowFlightMTLSInvalidCert into a single file. Just add catalogs to the query runner with those configurations, and then execute the test queries against those catalogs instead. This way we don't need to test 3 separate classes.

ZacBlanco · 2025-05-29T18:57:44Z

...-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/BaseArrowFlightClientHandler.java

        }
        catch (Exception e) {
-            throw new ArrowException(ARROW_FLIGHT_CLIENT_ERROR, "Error creating flight client: " + e.getMessage(), e);
+            Optional<Throwable> cause = Optional.ofNullable(e.getCause());


This optional also doesn't really provide value. There is no real use of the option value. Let's just use e.getCause() directly and do the cause instanceof <XYZ> check in the conditional instead.

Refactored to use e.getCause() instanceof ..

elbinpallimalilibm · 2025-05-30T01:52:22Z

One comment. Also, I think instead of having 3 test classes, let's combine the TestArrowFlightMTLSFails and TestArrowFlightMTLSInvalidCert into a single file. Just add catalogs to the query runner with those configurations, and then execute the test queries against those catalogs instead. This way we don't need to test 3 separate classes.

Merged the three test classes into one.

BryanCutler

Looks good, just a few minor nits

BryanCutler · 2025-06-02T21:33:37Z

presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowFlightConfig.java

+        return flightClientSSLCertificate;
+    }
+
+    @Config("arrow-flight.client-ssl-certificate")


instead of a javadoc, could you put this in a @ConfigDescription?

BryanCutler · 2025-06-02T21:33:41Z

presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowFlightConfig.java

+    }
+
+    /***
+     * Set the client SSL key used for mTLS authentication with Flight server


nit: period

same here about @ConfigDescription

presto-docs/src/main/sphinx/connector/base-arrow-flight.rst

BryanCutler

LGTM

elbinpallimalilibm requested review from a team, elharo and steveburnett as code owners May 23, 2025 02:50

elbinpallimalilibm requested a review from hantangwangd May 23, 2025 02:50

prestodb-ci added the from:IBM PR from IBM label May 23, 2025

prestodb-ci requested review from a team, bibith4 and pratyakshsharma and removed request for a team May 23, 2025 02:50

elbinpallimalilibm force-pushed the arrow_mtls branch from 38df200 to f2aff7e Compare May 23, 2025 02:58

elbinpallimalilibm marked this pull request as draft May 23, 2025 02:58

elbinpallimalilibm force-pushed the arrow_mtls branch from f2aff7e to 16c6399 Compare May 23, 2025 04:23

elbinpallimalilibm force-pushed the arrow_mtls branch 2 times, most recently from 7bdfc96 to bb2cf1d Compare May 24, 2025 00:10

elbinpallimalilibm marked this pull request as ready for review May 24, 2025 02:28

prestodb-ci requested a review from a team May 24, 2025 02:28

pratyakshsharma requested changes May 27, 2025

View reviewed changes

elbinpallimalilibm force-pushed the arrow_mtls branch from bb2cf1d to 60eaa3b Compare May 27, 2025 10:00

elbinpallimalilibm requested a review from pratyakshsharma May 27, 2025 10:19

elbinpallimalilibm force-pushed the arrow_mtls branch 2 times, most recently from 117c43a to 4a84c3b Compare May 27, 2025 11:55

steveburnett previously approved these changes May 27, 2025

View reviewed changes

elbinpallimalilibm force-pushed the arrow_mtls branch from 4a84c3b to 83a005f Compare May 27, 2025 14:07

pratyakshsharma reviewed May 28, 2025

View reviewed changes

elbinpallimalilibm dismissed steveburnett’s stale review via 1021d94 May 28, 2025 09:47

elbinpallimalilibm force-pushed the arrow_mtls branch from 83a005f to 1021d94 Compare May 28, 2025 09:47

elbinpallimalilibm requested a review from pratyakshsharma May 28, 2025 11:13

pratyakshsharma reviewed May 28, 2025

View reviewed changes

elbinpallimalilibm force-pushed the arrow_mtls branch 2 times, most recently from 3fe0e0e to 4bbfaa5 Compare May 28, 2025 13:20

pratyakshsharma reviewed May 28, 2025

View reviewed changes

elbinpallimalilibm force-pushed the arrow_mtls branch from 4bbfaa5 to dd0d61a Compare May 28, 2025 13:41

pratyakshsharma previously approved these changes May 28, 2025

View reviewed changes

ZacBlanco requested changes May 29, 2025

View reviewed changes

elbinpallimalilibm dismissed pratyakshsharma’s stale review via cd6d803 May 30, 2025 01:46

elbinpallimalilibm force-pushed the arrow_mtls branch from dd0d61a to cd6d803 Compare May 30, 2025 01:46

elbinpallimalilibm requested a review from ZacBlanco May 30, 2025 01:56

ZacBlanco previously approved these changes May 30, 2025

View reviewed changes

hantangwangd previously approved these changes Jun 1, 2025

View reviewed changes

ZacBlanco requested a review from BryanCutler June 2, 2025 04:22

BryanCutler reviewed Jun 2, 2025

View reviewed changes

elbinpallimalilibm dismissed stale reviews from hantangwangd and ZacBlanco via 2f38905 June 3, 2025 08:06

Add support for mTLS authentication in Arrow Flight client

e8a1c06

elbinpallimalilibm force-pushed the arrow_mtls branch from ad08755 to e8a1c06 Compare June 3, 2025 08:27

BryanCutler approved these changes Jun 3, 2025

View reviewed changes

ZacBlanco approved these changes Jun 3, 2025

View reviewed changes

ZacBlanco merged commit 8a04ee3 into prestodb:master Jun 3, 2025
99 checks passed

unidevel mentioned this pull request Jun 19, 2025

Add release notes for 0.294 unix280/presto#33

Closed

6 tasks

This was referenced Jul 4, 2025

Add release notes for 0.294 unix280/presto#35

Merged

Add release notes for 0.294 unix280/presto#36

Closed

Add release notes for 0.294 unix280/presto#37

Closed

This was referenced Jul 24, 2025

Add release notes for 0.294 unix280/presto#39

Merged

Add release notes for 0.294 unix280/presto#40

Merged

prestodb-ci mentioned this pull request Jul 28, 2025

Add release notes for 0.294 #25633

Merged

6 tasks


		abstract Map<String, String> getCatalogProperties();

		protected int getServerPort()

	if (cause.filter(c -> c instanceof InvalidKeyException).isPresent()) {
	if (e instanceOf InvalidKeyException) {

	logger.error("Error closing input stream", e);
	logger.error("Error closing input stream for clientCertificate", e);

Conversation

elbinpallimalilibm commented May 23, 2025

Description

Motivation and Context

Impact

Test Plan

Contributor checklist

Release Notes

Uh oh!

elbinpallimalilibm commented May 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

steveburnett commented May 23, 2025

Uh oh!

pratyakshsharma left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

steveburnett left a comment

Choose a reason for hiding this comment

Uh oh!

pratyakshsharma left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pratyakshsharma left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pratyakshsharma left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

elbinpallimalilibm commented May 23, 2025 •

edited

Loading