fix: stats for single shard (#992)

Author: ferhatelmas

src/internal/sharding/sharder.ts

@@ -1,4 +1,4 @@
-import { ResourceKind, ShardRow, ShardStatus } from './store'
+import { ResourceKind, ShardRow, ShardStats, ShardStatus } from './store'
 
 export interface ShardResource {
   kind: ResourceKind
@@ -7,6 +7,14 @@ export interface ShardResource {
   logicalName: string
 }
 
+export interface ReservationResult {
+  reservationId: string
+  shardId: string
+  shardKey: string
+  slotNo: number
+  leaseExpiresAt: string
+}
+
 export interface Sharder {
   createShard(opts: {
     kind: ResourceKind
@@ -17,26 +25,13 @@ export interface Sharder {
 
   setShardStatus(shardId: string | number, status: ShardStatus): Promise<void>
 
-  reserve(
-    opts: ShardResource & {
-      kind: ResourceKind
-      tenantId: string
-      bucketName: string
-      logicalName: string
-    }
-  ): Promise<{
-    reservationId: string
-    shardId: string
-    shardKey: string
-    slotNo: number
-    leaseExpiresAt: string
-  }>
+  reserve(opts: ShardResource): Promise<ReservationResult>
   confirm(reservationId: string, resource: ShardResource): Promise<void>
   cancel(reservationId: string): Promise<void>
   expireLeases(): Promise<number>
   freeByLocation(shardId: string | number, slotNo: number): Promise<void>
   freeByResource(shardId: string | number, resource: ShardResource): Promise<void>
-  shardStats(kind?: ResourceKind): Promise<any>
+  shardStats(kind?: ResourceKind): Promise<ShardStats>
   findShardByResourceId(param: ShardResource): Promise<ShardRow | null>
   listShardByKind(icebergTables: ResourceKind): Promise<ShardRow[]>
 

src/internal/sharding/store.ts

@@ -25,6 +25,14 @@ export type ReservationRow = {
   created_at: string
 }
 
+export type ShardStats = Array<{
+  shardId: string
+  shardKey: string
+  capacity: number
+  used: number
+  free: number
+}>
+
 /** Factory that opens a transaction and passes a store bound to that tx */
 export interface ShardStoreFactory<Tnx = unknown> {
   withTransaction<T>(fn: (store: ShardStore) => Promise<T>): Promise<T>
@@ -101,11 +109,7 @@ export interface ShardStore {
   findShardByResourceId(tenantId: string, resourceId: string): Promise<ShardRow | null>
 
   // Stats
-  shardStats(
-    kind?: ResourceKind
-  ): Promise<
-    Array<{ shardId: string; shardKey: string; capacity: number; used: number; free: number }>
-  >
+  shardStats(kind?: ResourceKind): Promise<ShardStats>
 
   findShardById(shardId: number): Promise<ShardRow | null>
 }

src/internal/sharding/strategy/single-shard.ts

@@ -1,5 +1,5 @@
-import { ResourceKind, ShardRow, ShardStatus } from '@internal/sharding/store'
-import { Sharder, ShardResource } from '../sharder'
+import { ResourceKind, ShardRow, ShardStats, ShardStatus } from '@internal/sharding/store'
+import { ReservationResult, Sharder, ShardResource } from '../sharder'
 
 export class SingleShard implements Sharder {
   constructor(
@@ -9,7 +9,7 @@ export class SingleShard implements Sharder {
     }
   ) {}
 
-  listShardByKind(): Promise<ShardRow[]> {
+  listShardByKind(_kind: ResourceKind): Promise<ShardRow[]> {
     return Promise.resolve([
       {
         id: 1,
@@ -23,32 +23,34 @@ export class SingleShard implements Sharder {
     ])
   }
 
-  shardStats(): Promise<any> {
-    return Promise.resolve({
-      shardId: 1,
-      shardKey: this.singleShard.shardKey,
-      capacity: this.singleShard.capacity,
-      used: -1,
-      free: -1,
-    })
+  shardStats(_kind?: ResourceKind): Promise<ShardStats> {
+    return Promise.resolve([
+      {
+        shardId: '1',
+        shardKey: this.singleShard.shardKey,
+        capacity: this.singleShard.capacity,
+        used: -1,
+        free: -1,
+      },
+    ])
   }
 
-  withTnx(): Sharder {
+  withTnx(_tnx: unknown): Sharder {
     return new SingleShard({
       shardKey: this.singleShard.shardKey,
       capacity: this.singleShard.capacity,
     })
   }
 
-  freeByResource(): Promise<void> {
+  freeByResource(_shardId: string | number, _resource: ShardResource): Promise<void> {
     return Promise.resolve()
   }
 
-  cancel(): Promise<void> {
+  cancel(_reservationId: string): Promise<void> {
     return Promise.resolve(undefined)
   }
 
-  confirm(): Promise<void> {
+  confirm(_reservationId: string, _resource: ShardResource): Promise<void> {
     return Promise.resolve(undefined)
   }
 
@@ -85,31 +87,21 @@ export class SingleShard implements Sharder {
     })
   }
 
-  freeByLocation(): Promise<void> {
+  freeByLocation(_shardId: string | number, _slotNo: number): Promise<void> {
     return Promise.resolve(undefined)
   }
 
-  reserve(): Promise<{
-    reservationId: string
-    shardId: string
-    shardKey: string
-    slotNo: number
-    leaseExpiresAt: string
-  }> {
+  reserve(_opts: ShardResource): Promise<ReservationResult> {
     return Promise.resolve({
       leaseExpiresAt: '',
       reservationId: '',
-      shardId: this.singleShard.shardKey,
+      shardId: '1',
       shardKey: this.singleShard.shardKey,
       slotNo: 0,
     })
   }
 
-  setShardStatus(): Promise<void> {
-    return Promise.resolve(undefined)
-  }
-
-  shardStatsByKind(): Promise<any> {
+  setShardStatus(_shardId: string | number, _status: ShardStatus): Promise<void> {
     return Promise.resolve(undefined)
   }
 }

src/test/sharding.test.ts

@@ -2,7 +2,7 @@
 
 import { multitenantKnex } from '@internal/database'
 import { runMultitenantMigrations } from '@internal/database/migrations'
-import { KnexShardStoreFactory, ShardCatalog } from '@internal/sharding'
+import { KnexShardStoreFactory, ShardCatalog, SingleShard } from '@internal/sharding'
 import {
   ExpiredReservationError,
   NoActiveShardError,
@@ -798,3 +798,22 @@ describe('Sharding System', () => {
     })
   })
 })
+
+describe('SingleShard', () => {
+  it('returns shard stats in the canonical array shape', async () => {
+    const sharder = new SingleShard({
+      shardKey: 'single-shard-key',
+      capacity: 25,
+    })
+
+    await expect(sharder.shardStats()).resolves.toEqual([
+      {
+        shardId: '1',
+        shardKey: 'single-shard-key',
+        capacity: 25,
+        used: -1,
+        free: -1,
+      },
+    ])
+  })
+})