feat: add support for setting the default database (#9)

Some flight SQL servers support the metadata field `database` to
indicate the default database that queries should use (influx and
lancedb enterprise for example). This PR adds support for that option.
Regrettably, we cannot test this with integration tests because the
voltron flight server does not support this field.

Author: westonpace

README.md

@@ -12,15 +12,15 @@ Currently all testing is done on Node.
 You can install the client using npm:
 
 ```bash
-npm install @lancedb/flightsql-client
+npm install @lancedb/arrow-flight-sql-client
 ```
 
 ## Usage
 
 To use the client, you first need to connect to your database:
 
 ```javascript
-import { Client } from "@lancedb/flightsql-client";
+import { Client } from "@lancedb/arrow-flight-sql-client";
 
 const client = await Client.connect({
   host: "mydb.com:10025",

src/client.ts

@@ -20,6 +20,10 @@ export interface ClientOptions {
    * The password to login with
    */
   password: string;
+  /**
+   * The default database to use (non-standard, only supported on some servers)
+   */
+  defaultDatabase?: string;
   /**
    * Is the server using TLS?  If not, this must be set to true.
    */
@@ -118,7 +122,12 @@ export class Client {
    * @returns A client that can be used to execute queries
    */
   public static async connect(options: ClientOptions): Promise<Client> {
-    const sql = await FlightSqlClient.connect(options.host, options.username, options.password);
+    const sql = await FlightSqlClient.connect(
+      options.host,
+      options.username,
+      options.password,
+      options.defaultDatabase,
+    );
     return new Client(sql);
   }
 

src/flightsql.ts

@@ -29,7 +29,7 @@ export class FlightSqlClient {
     this.flight = flight;
   }
 
-  private async login(username: string, password: string): Promise<void> {
+  private async login(username: string, password: string, defaultDatabase?: string): Promise<void> {
     // Most servers seem to use Basic auth for the handshake to get a token.
     // The GRPC headers (and not the handshake request / response) are used to
     // transmit the username, password, and token.
@@ -50,23 +50,26 @@ export class FlightSqlClient {
       // The logic here is maybe not ideal.  If a server someone sends an empty metadata message and then
       // a populated data payload message, then we will see the empty payload message first and fail to
       // wait for the next message.  We can optimize this later if we see it in the wild.
-      const hello_rsp = await firstValueFrom(call.responses);
+      const helloRsp = await firstValueFrom(call.responses);
 
-      if (hello_rsp?.data?.payload) {
+      const defaultMetadata = {};
+      if (defaultDatabase) {
+        defaultMetadata["database"] = defaultDatabase;
+      }
+
+      if (helloRsp?.data?.payload) {
         // If we get a payload prefer that as a token
-        const payload = new TextDecoder().decode(hello_rsp.data.payload);
-        this.flight.set_default_metadata({
-          authorization: "Bearer " + payload,
-        });
-      } else if (hello_rsp?.metadata) {
+        const payload = new TextDecoder().decode(helloRsp.data.payload);
+        defaultMetadata["authorization"] = "Bearer " + payload;
+        this.flight.set_default_metadata(defaultMetadata);
+      } else if (helloRsp?.metadata) {
         // Otherwise if we get metadata then use that as the token
-        const authorization = hello_rsp.metadata.get_first_string("authorization");
+        const authorization = helloRsp.metadata.get_first_string("authorization");
         if (!authorization) {
           throw new Error("Handshake failed, metadata received but no authorization header present");
         }
-        this.flight.set_default_metadata({
-          authorization,
-        });
+        defaultMetadata["authorization"] = authorization;
+        this.flight.set_default_metadata(defaultMetadata);
       } else {
         throw new Error("Handshake failed, no metadata or data received from server before call completed");
       }
@@ -91,10 +94,15 @@ export class FlightSqlClient {
    * @param password The password to use for the handshake
    * @returns A client that can be used to execute queries
    */
-  public static async connect(host: string, username: string, password: string): Promise<FlightSqlClient> {
+  public static async connect(
+    host: string,
+    username: string,
+    password: string,
+    defaultDatabase?: string,
+  ): Promise<FlightSqlClient> {
     const sql = new FlightClient(host);
     const client = new FlightSqlClient(sql);
-    await client.login(username, password);
+    await client.login(username, password, defaultDatabase);
     return client;
   }