Add type safety to the limit funcion

tursodatabase · Jun 21, 2024 · fab870e · fab870e
1 parent 6bf5753
commit fab870e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 12 deletions.
diff --git a/packages/libsql-client/src/http.ts b/packages/libsql-client/src/http.ts
@@ -79,7 +79,7 @@ export class HttpClient implements Client {
     #url: URL;
     #authToken: string | undefined;
     #isSchemaDatabase: Promise<boolean> | undefined;
-    #limit: ReturnType<typeof promiseLimit<any>>;
+    #promiseLimitFunction: ReturnType<typeof promiseLimit<any>>;
 
     /** @private */
     constructor(
@@ -94,7 +94,7 @@ export class HttpClient implements Client {
         this.protocol = "http";
         this.#url = url;
         this.#authToken = authToken;
-        this.#limit = promiseLimit<any>(concurrency);
+        this.#promiseLimitFunction = promiseLimit<any>(concurrency);
     }
 
     getIsSchemaDatabase(): Promise<boolean> {
@@ -108,8 +108,12 @@ export class HttpClient implements Client {
         return this.#isSchemaDatabase;
     }
 
+    private async limit<T>(fn: () => Promise<T>): Promise<T> {
+        return this.#promiseLimitFunction(fn);
+    }
+
     async execute(stmt: InStatement): Promise<ResultSet> {
-        return this.#limit(async () => {
+        return this.limit<ResultSet>(async () => {
             try {
                 const isSchemaDatabasePromise = this.getIsSchemaDatabase();
                 const hranaStmt = stmtToHrana(stmt);
@@ -144,7 +148,7 @@ export class HttpClient implements Client {
         stmts: Array<InStatement>,
         mode: TransactionMode = "deferred",
     ): Promise<Array<ResultSet>> {
-        return this.#limit(async () => {
+        return this.limit<Array<ResultSet>>(async () => {
             try {
                 const isSchemaDatabasePromise = this.getIsSchemaDatabase();
                 const hranaStmts = stmts.map(stmtToHrana);
@@ -194,7 +198,7 @@ export class HttpClient implements Client {
     async transaction(
         mode: TransactionMode = "write",
     ): Promise<HttpTransaction> {
-        return this.#limit(async () => {
+        return this.limit<HttpTransaction>(async () => {
             try {
                 const version = await this.#client.getVersion();
                 return new HttpTransaction(
@@ -209,7 +213,7 @@ export class HttpClient implements Client {
     }
 
     async executeMultiple(sql: string): Promise<void> {
-        this.#limit(async () => {
+        this.limit<void>(async () => {
             try {
                 // Pipeline all operations, so `hrana.HttpClient` can open the stream, execute the sequence and
                 // close the stream in a single HTTP request.

diff --git a/packages/libsql-client/src/ws.ts b/packages/libsql-client/src/ws.ts
@@ -132,7 +132,7 @@ export class WsClient implements Client {
     closed: boolean;
     protocol: "ws";
     #isSchemaDatabase: Promise<boolean> | undefined;
-    #limit: ReturnType<typeof promiseLimit<any>>;
+    #promiseLimitFunction: ReturnType<typeof promiseLimit<any>>;
 
     /** @private */
     constructor(
@@ -149,7 +149,7 @@ export class WsClient implements Client {
         this.#futureConnState = undefined;
         this.closed = false;
         this.protocol = "ws";
-        this.#limit = promiseLimit(concurrency);
+        this.#promiseLimitFunction = promiseLimit<any>(concurrency);
     }
 
     getIsSchemaDatabase(): Promise<boolean> {
@@ -163,8 +163,12 @@ export class WsClient implements Client {
         return this.#isSchemaDatabase;
     }
 
+    private async limit<T>(fn: () => Promise<T>): Promise<T> {
+        return this.#promiseLimitFunction(fn);
+    }
+
     async execute(stmt: InStatement): Promise<ResultSet> {
-        return this.#limit(async () => {
+        return this.limit<ResultSet>(async () => {
             const streamState = await this.#openStream();
             try {
                 const isSchemaDatabasePromise = this.getIsSchemaDatabase();
@@ -198,7 +202,7 @@ export class WsClient implements Client {
         stmts: Array<InStatement>,
         mode: TransactionMode = "deferred",
     ): Promise<Array<ResultSet>> {
-        return this.#limit(async () => {
+        return this.limit<Array<ResultSet>>(async () => {
             const streamState = await this.#openStream();
             try {
                 const isSchemaDatabasePromise = this.getIsSchemaDatabase();
@@ -235,7 +239,7 @@ export class WsClient implements Client {
     }
 
     async transaction(mode: TransactionMode = "write"): Promise<WsTransaction> {
-        return this.#limit(async () => {
+        return this.limit<WsTransaction>(async () => {
             const streamState = await this.#openStream();
             try {
                 const version = await streamState.conn.client.getVersion();
@@ -250,7 +254,7 @@ export class WsClient implements Client {
     }
 
     async executeMultiple(sql: string): Promise<void> {
-        this.#limit(async () => {
+        this.limit<void>(async () => {
             const streamState = await this.#openStream();
             try {
                 // Schedule all operations synchronously, so they will be pipelined and executed in a single