Allow to store externalAuthenticationToken in SYSTEM Cache

client/trino-client/src/main/java/io/trino/client/auth/external/KnownToken.java

@@ -31,4 +31,9 @@ static KnownToken memoryCached()
     {
         return MemoryCachedKnownToken.INSTANCE;
     }
+
+    static KnownToken systemCached()
+    {
+        return SystemCachedKnownToken.INSTANCE;
+    }
 }

client/trino-client/src/main/java/io/trino/client/auth/external/SystemCachedKnownToken.java

@@ -0,0 +1,151 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.client.auth.external;
+
+import com.google.common.collect.ImmutableSet;
+import dev.failsafe.Failsafe;
+import dev.failsafe.RetryPolicy;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.Optional;
+import java.util.function.Supplier;
+
+import static java.nio.file.attribute.PosixFilePermission.OWNER_READ;
+import static java.nio.file.attribute.PosixFilePermission.OWNER_WRITE;
+import static java.time.temporal.ChronoUnit.MILLIS;
+import static java.util.Objects.requireNonNull;
+
+/**
+ * This KnownToken instance persists the token to ~/.trino/.token on the filesystem,
+ * allowing it to be reused across separate CLI invocations.
+ * A lock file (~/.trino/.token.lck) is used to coordinate token acquisition
+ * across processes — its atomic creation acts as a cross-process tryLock.
+ * The implementation is similar to MemoryCachedKnownToken, but with LOCK_FILE acting as a Lock
+ */
+class SystemCachedKnownToken
+        implements KnownToken
+{
+    private static final Path DEFAULT_TRINO_DIR = Path.of(System.getProperty("user.home"), ".trino");
+    // duration for the user to authenticate within IDP. It involves clicking and typing from a real person within a browser, so should be counted in minutes.
+    private static final Duration DEFAULT_LOCK_MAX_WAIT = Duration.ofMinutes(10);
+
+    public static final SystemCachedKnownToken INSTANCE = new SystemCachedKnownToken(DEFAULT_TRINO_DIR);
+
+    private final Path trinoDir;
+    private final Path tokenFile;
+    private final Path lockFile;
+    private final Duration lockMaxWait;
+
+    SystemCachedKnownToken(Path trinoDir)
+    {
+        this(trinoDir, DEFAULT_LOCK_MAX_WAIT);
+    }
+
+    SystemCachedKnownToken(Path trinoDir, Duration lockMaxWait)
+    {
+        this.trinoDir = requireNonNull(trinoDir, "trinoDir is null");
+        this.tokenFile = trinoDir.resolve(".token");
+        this.lockFile = trinoDir.resolve(".token.lck");
+        this.lockMaxWait = requireNonNull(lockMaxWait, "lockMaxWait is null");
+    }
+
+    @Override
+    public Optional<Token> getToken()
+    {
+        // Wait while lock file exists, mimicking readLock blocking while writeLock is held
+        boolean lockFileExists = Failsafe.with(RetryPolicy.<Boolean>builder()
+                        .handleResultIf(Boolean.TRUE::equals)
+                        .withMaxAttempts(-1)
+                        .withDelay(100, 1000, MILLIS)
+                        .withMaxDuration(lockMaxWait)
+                        .build())
+                .get(() -> Files.exists(lockFile));
+        if (lockFileExists) {
+            throw new IllegalStateException("Lock file " + lockFile + " for System Cached token still exists after waiting " + lockMaxWait + ". " +
+                    "It may be created by another concurrent authentication, which is still in progress. " +
+                    "If it's not a case and another transaction was abandoned - please remove the lock file manually and retry authentication.");
+        }
+
+        if (!Files.exists(tokenFile)) {
+            return Optional.empty();
+        }
+        try {
+            String content = Files.readString(tokenFile).trim();
+            if (content.isEmpty()) {
+                return Optional.empty();
+            }
+            return Optional.of(new Token(content));
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException("Failed to read token from " + tokenFile, e);
+        }
+    }
+
+    @Override
+    public void setupToken(Supplier<Optional<Token>> tokenSource)
+    {
+        try {
+            Files.createDirectories(trinoDir);
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException("Failed to create directory " + trinoDir, e);
+        }
+
+        // Atomically create the lock file. If it already exists, another process
+        // is obtaining a token — skip, just like MemoryCachedKnownToken's tryLock.
+        try {
+            Files.createFile(lockFile);
+        }
+        catch (FileAlreadyExistsException e) {
+            return;
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException("Failed to create lock file " + lockFile, e);
+        }
+
+        try {
+            // Clear token before obtaining new one, as it might fail leaving old invalid token.
+            Files.deleteIfExists(tokenFile);
+            Optional<Token> token = tokenSource.get();
+            token.ifPresent(this::writeTokenToFile);
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException("Failed to update token file " + tokenFile, e);
+        }
+        finally {
+            try {
+                Files.deleteIfExists(lockFile);
+            }
+            catch (IOException e) {
+                throw new UncheckedIOException("Failed to delete lock file " + lockFile, e);
+            }
+        }
+    }
+
+    private void writeTokenToFile(Token token)
+    {
+        try {
+            Files.writeString(tokenFile, token.token());
+            Files.setPosixFilePermissions(tokenFile, ImmutableSet.of(OWNER_READ, OWNER_WRITE));
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException("Failed to write token to " + tokenFile, e);
+        }
+    }
+}

client/trino-client/src/main/java/io/trino/client/uri/KnownTokenCache.java

@@ -30,6 +30,13 @@ KnownToken create()
         {
             return KnownToken.memoryCached();
         }
+    },
+    SYSTEM {
+        @Override
+        KnownToken create()
+        {
+            return KnownToken.systemCached();
+        }
     };
 
     abstract KnownToken create();

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

@@ -25,10 +25,12 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.api.io.TempDir;
 import org.junit.jupiter.api.parallel.Execution;
 
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.nio.file.Path;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
@@ -193,16 +195,33 @@ public void testAuthenticationFromMultipleThreadsWithLocallyStoredToken()
 
     @Test
     @Timeout(2)
-    public void testAuthenticationFromMultipleThreadsWithCachedToken()
+    public void testAuthenticationFromMultipleThreadsWithMemoryCachedToken()
+    {
+        KnownToken knownToken = KnownToken.memoryCached();
+        testAuthenticationFromMultipleThreadsWithCachedToken(knownToken);
+    }
+
+    @Test
+    @Timeout(10)
+    public void testAuthenticationFromMultipleThreadsWithSystemCachedToken(@TempDir Path tempDir)
+    {
+        KnownToken knownToken = new SystemCachedKnownToken(tempDir);
+        testAuthenticationFromMultipleThreadsWithCachedToken(knownToken);
+    }
+
+    private static void testAuthenticationFromMultipleThreadsWithCachedToken(KnownToken knownToken)
     {
         MockTokenPoller tokenPoller = new MockTokenPoller()
-                .withResult(URI.create("http://token.uri"), successful(new Token("valid-token")));
+                // will be cached and reused
+                .withResult(URI.create("http://token.uri"), successful(new Token("first-token")))
+                // should never be emitted
+                .withResult(URI.create("http://token.uri"), successful(new Token("second-token")));
         MockRedirectHandler redirectHandler = new MockRedirectHandler()
                 .sleepOnRedirect(Duration.ofSeconds(1));
 
         List<Future<Request>> requests = times(
-                2,
-                () -> new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1))
+                100,
+                () -> new ExternalAuthenticator(redirectHandler, tokenPoller, knownToken, Duration.ofSeconds(1))
                         .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\"")))
                 .map(executor::submit)
                 .collect(toImmutableList());
@@ -211,7 +230,7 @@ public void testAuthenticationFromMultipleThreadsWithCachedToken()
         assertion.requests()
                 .extracting(Request::headers)
                 .extracting(headers -> headers.get(AUTHORIZATION))
-                .containsOnly("Bearer valid-token");
+                .containsOnly("Bearer first-token");
         assertion.assertThatNoExceptionsHasBeenThrown();
         assertThat(redirectHandler.getRedirectionCount()).isEqualTo(1);
     }