diff --git a/README.md b/README.md index 0c2e751..fe1ca3f 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ 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 @@ -20,7 +20,7 @@ npm install @lancedb/flightsql-client 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", diff --git a/src/client.ts b/src/client.ts index 67753f9..afa4eca 100644 --- a/src/client.ts +++ b/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 { - 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); } diff --git a/src/flightsql.ts b/src/flightsql.ts index 997130f..7a58cf0 100644 --- a/src/flightsql.ts +++ b/src/flightsql.ts @@ -29,7 +29,7 @@ export class FlightSqlClient { this.flight = flight; } - private async login(username: string, password: string): Promise { + private async login(username: string, password: string, defaultDatabase?: string): Promise { // 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 { + public static async connect( + host: string, + username: string, + password: string, + defaultDatabase?: string, + ): Promise { const sql = new FlightClient(host); const client = new FlightSqlClient(sql); - await client.login(username, password); + await client.login(username, password, defaultDatabase); return client; }