Support userAgent and make LookerResponseParser public for testing

build.gradle.kts

@@ -177,7 +177,7 @@ allprojects {
         dependencies {
             "implementation"(platform(project(":bom")))
             // Add the locally bundled LookerSDK fat jar
-            "implementation"(files("../libs/looker-kotlin-sdk-f91f8ad.jar"))
+            "implementation"(files("../libs/looker-kotlin-sdk-a48011f.jar"))
         }
     }
     if (!skipAutostyle) {

core/src/main/java/org/apache/calcite/avatica/remote/looker/LookerDriver.java

@@ -17,6 +17,7 @@
 package org.apache.calcite.avatica.remote.looker;
 
 import org.apache.calcite.avatica.AvaticaConnection;
+import org.apache.calcite.avatica.ConnectStringParser;
 import org.apache.calcite.avatica.DriverVersion;
 import org.apache.calcite.avatica.Meta;
 import org.apache.calcite.avatica.UnregisteredDriver;
@@ -83,8 +84,11 @@ public Meta createMeta(AvaticaConnection connection) {
     Service service = conn.getService();
     assert service instanceof LookerRemoteService;
 
+    // puts all additional url params into properties
+    Properties properties = ConnectStringParser.parse(url, info);
+
     // create and set LookerSDK for the connection
-    LookerSDK sdk = LookerSdkFactory.createSdk(conn.config().url(), info);
+    LookerSDK sdk = LookerSdkFactory.createSdk(conn.config().url(), properties);
     ((LookerRemoteService) service).setSdk(sdk);
     return conn;
   }

core/src/main/java/org/apache/calcite/avatica/remote/looker/LookerFrameEnvelope.java

@@ -23,16 +23,24 @@
  * {@link LookerRemoteMeta#stmtQueueMap} to hold complete Frames and present exceptions when
  * consumers are ready to encounter them.
  */
-class LookerFrameEnvelope {
+public class LookerFrameEnvelope {
 
-  final Frame frame;
-  final Exception exception;
+  private final Frame frame;
+  private final Exception exception;
 
   private LookerFrameEnvelope(/*@Nullable*/ Frame frame, /*@Nullable*/ Exception exception) {
     this.frame = frame;
     this.exception = exception;
   }
 
+  public Frame getFrame() {
+    return this.frame;
+  }
+
+  public Exception getException() {
+    return this.exception;
+  }
+
   /**
    * Constructs a LookerFrameEnvelope with a {@link Frame}.
    */

core/src/main/java/org/apache/calcite/avatica/remote/looker/LookerRemoteMeta.java

@@ -111,7 +111,7 @@ static class LookerFrame extends Frame {
     }
 
     /**
-     * Creates a {@code LookerFrame} for the statement slug
+     * Creates a {@code LookerFrame} for the query id.
      *
      * @param sqlInterfaceQueryId id for the prepared statement generated by a Looker instance.
      * @return the {@code firstFrame} for the result set.
@@ -176,17 +176,22 @@ protected InputStream makeRunQueryRequest(String url) throws IOException, SQLExc
       trustAllHosts(connection);
     }
 
-    // timeout is given as seconds
+    // convert timeout seconds to milliseconds
     int timeout = sdkTransport.getOptions().getTimeout() * 1000;
     connection.setReadTimeout(timeout);
+
     connection.setRequestMethod("GET");
     connection.setRequestProperty("Accept", "application/json");
     connection.setDoOutput(true);
 
-    // Set the auth header as the SDK would
+    // set the auth header as the SDK would
     connection.setRequestProperty("Authorization",
         "token " + authSession.getAuthToken().getAccessToken());
 
+    // copy the headers from the authenticated SDK.
+    authSession.getApiSettings().getHeaders()
+        .forEach((header, value) -> connection.setRequestProperty(header, value));
+
     int responseCode = connection.getResponseCode();
     if (responseCode == 200) {
       // return the input stream to parse from.
@@ -198,7 +203,7 @@ protected InputStream makeRunQueryRequest(String url) throws IOException, SQLExc
               .lines()
               .collect(Collectors.joining("\n"));
       HashMap<String, String> errorMap = JsonService.MAPPER.readValue(errorMessage, HashMap.class);
-      throw new SQLException("Looker generated SQL failed to execute: "
+      throw new SQLException("Looker generated SQL failed to execute. "
           + errorMap.getOrDefault("message", ""));
     }
   }
@@ -229,11 +234,11 @@ public Frame fetch(final StatementHandle h, final long offset, final int fetchMa
         // remove the statement from the map if it has an exception, or it is the last frame
         if (nextEnvelope.hasException()) {
           stmtQueueMap.remove(h.id);
-          throw new RuntimeException(nextEnvelope.exception);
-        } else if (nextEnvelope.frame.done) {
+          throw new RuntimeException(nextEnvelope.getException());
+        } else if (nextEnvelope.getFrame().done) {
           stmtQueueMap.remove(h.id);
         }
-        return nextEnvelope.frame;
+        return nextEnvelope.getFrame();
       } catch (InterruptedException e) {
         throw new RuntimeException(e);
       }

core/src/main/java/org/apache/calcite/avatica/remote/looker/LookerRemoteService.java

@@ -21,9 +21,9 @@
 import org.apache.calcite.avatica.remote.JsonService;
 import org.apache.calcite.avatica.remote.looker.LookerRemoteMeta.LookerFrame;
 
-import com.looker.sdk.JdbcInterface;
 import com.looker.sdk.LookerSDK;
 import com.looker.sdk.SqlInterfaceQuery;
+import com.looker.sdk.SqlInterfaceQueryMetadata;
 import com.looker.sdk.WriteSqlInterfaceQueryCreate;
 
 import java.io.IOException;
@@ -108,13 +108,13 @@ private SqlQueryWithSignature prepareQuery(String sql) {
   /**
    * Handles all non-overridden {@code apply} methods.
    *
-   * Calls the {@code jdbc_interface} endpoint of the instance which behaves similarly to a standard
-   * Avatica server.
+   * Calls the {@code sql_interface_metadata} endpoint of the instance which behaves similarly to a
+   * standard Avatica server.
    */
   @Override
   public String apply(String request) {
     checkSdk();
-    JdbcInterface response = safeSdkCall(() -> sdk.jdbc_interface(request));
+    SqlInterfaceQueryMetadata response = safeSdkCall(() -> sdk.sql_interface_metadata(request));
 
     return response.getResults();
   }

core/src/main/java/org/apache/calcite/avatica/remote/looker/LookerResponseParser.java

@@ -43,7 +43,7 @@ public class LookerResponseParser {
 
   private final BlockingQueue<LookerFrameEnvelope> queue;
 
-  LookerResponseParser(BlockingQueue<LookerFrameEnvelope> queue) {
+  public LookerResponseParser(BlockingQueue<LookerFrameEnvelope> queue) {
     assert queue != null : "null queue!";
 
     this.queue = queue;

core/src/main/java/org/apache/calcite/avatica/remote/looker/LookerSdkFactory.java

@@ -17,6 +17,7 @@
 package org.apache.calcite.avatica.remote.looker;
 
 import com.looker.rtl.AuthSession;
+import com.looker.rtl.AuthToken;
 import com.looker.rtl.ConfigurationProvider;
 import com.looker.rtl.SDKErrorInfo;
 import com.looker.rtl.SDKResponse;
@@ -48,6 +49,11 @@ private LookerSdkFactory() {
   private static final String RESULT_FORMAT = "json_bi";
   private static final String QUERY_ENDPOINT = "/api/4.0/sql_interface_queries/%s/run/%s";
 
+  /**
+   * 1 hour in seconds. This is not configurable.
+   */
+  private static final Long SESSION_LENGTH = 3600L;
+
 
   /**
    * Simple interface to wrap SDK calls
@@ -106,26 +112,36 @@ private static AuthSession createAuthSession(String url, Map<String, String> pro
     if (apiLogin && authToken) {
       throw new SQLInvalidAuthorizationSpecException("Invalid connection params.\n"
           + "Cannot provide both API3 credentials and an access token");
-    } else if (apiLogin) {
-      apiConfig.put("client_id", props.get("user"));
-      apiConfig.put("client_secret", props.get("password"));
-    } else if (authToken) {
-      // TODO b/295025684: Set the token for the session using `session.setAuthToken(AuthToken);`.
-      //  Doing so will allow us to rely on the same auth session for the stream query call.
-      Map<String, String> headers = new HashMap<>();
-      headers.put("Authorization", "token " + props.get("token"));
-      apiConfig.put("headers", headers.toString());
-    } else {
+    }
+    if (!apiLogin && !authToken) {
       throw new SQLInvalidAuthorizationSpecException(
           "Invalid connection params.\n" + "Missing either API3 credentials or access token");
     }
 
+    if (apiLogin) {
+      apiConfig.put("client_id", props.get("user"));
+      apiConfig.put("client_secret", props.get("password"));
+    }
+
     ConfigurationProvider finalizedConfig = ApiSettings.fromMap(apiConfig);
+
+    String userAgent = props.get("userAgent");
+    if (userAgent != null) {
+      Map<String, String> headers = finalizedConfig.getHeaders();
+      headers.put("User-Agent", props.get("userAgent"));
+      finalizedConfig.setHeaders(headers);
+    }
+
     AuthSession session = new AuthSession(finalizedConfig, new Transport(finalizedConfig));
+
     // need to log in if client_id and client_secret are used
     if (apiLogin) {
       // empty string means no sudo - we won't support this
       session.login("");
+    } else if (authToken) {
+      // set the auth token if one was supplied from the OAuth flow
+      session.setAuthToken(
+          new AuthToken(props.get("token"), "Bearer", SESSION_LENGTH, null));
     }
 
     return session;

core/src/test/java/org/apache/calcite/avatica/remote/looker/LookerTestCommon.java

@@ -30,6 +30,8 @@ private LookerTestCommon() {
 
   private static final Properties BASE_PROPS = new Properties();
 
+  static final String USER_AGENT = ";userAgent=FooBarInc";
+
   static {
     BASE_PROPS.put("user", CLIENT_ID);
     BASE_PROPS.put("password", CLIENT_SECRET);