prisma-sharding 0.0.6 → 0.0.7

package/dist/index.mjs

@@ -96,6 +96,11 @@ var ShardRouter = class {
     if (this.strategy === "consistent-hash") {
       this.initializeConsistentHashRing();
     }
+    this.sortedRingHashes = Array.from(this.consistentHashRing.keys()).sort((a, b) => a - b);
+    this.randomShardWeights = this.shardIds.map(
+      (shardId) => config.shardWeights?.get(shardId) ?? 1
+    );
+    this.totalRandomWeight = this.randomShardWeights.reduce((sum, weight) => sum + weight, 0);
   }
   initializeConsistentHashRing() {
     for (const shardId of this.shardIds) {
@@ -120,25 +125,49 @@ var ShardRouter = class {
     return hash % shardCount;
   }
   getIndexConsistentHash(key) {
+    return this.shardIds.indexOf(this.getConsistentHashShardId(key));
+  }
+  getConsistentHashShardId(key) {
     const hash = hashString(key);
-    const sortedHashes = Array.from(this.consistentHashRing.keys()).sort((a, b) => a - b);
-    for (const ringHash of sortedHashes) {
-      if (hash <= ringHash) {
-        const shardId = this.consistentHashRing.get(ringHash);
-        const match2 = shardId.match(/shard_(\d+)/);
-        return match2 ? parseInt(match2[1], 10) - 1 : 0;
+    let low = 0;
+    let high = this.sortedRingHashes.length - 1;
+    let matchIndex = 0;
+    while (low <= high) {
+      const middle = Math.floor((low + high) / 2);
+      if (this.sortedRingHashes[middle] >= hash) {
+        matchIndex = middle;
+        high = middle - 1;
+      } else {
+        low = middle + 1;
       }
     }
-    const firstShardId = this.consistentHashRing.get(sortedHashes[0]);
-    const match = firstShardId.match(/shard_(\d+)/);
-    return match ? parseInt(match[1], 10) - 1 : 0;
+    if (low >= this.sortedRingHashes.length) {
+      matchIndex = 0;
+    }
+    return this.consistentHashRing.get(this.sortedRingHashes[matchIndex]);
   }
   getShardId(key) {
+    if (this.shardIds.length === 0) {
+      throw new RoutingError("No shards available");
+    }
+    if (this.strategy === "consistent-hash") {
+      return this.getConsistentHashShardId(key);
+    }
     const index = this.getShardIndex(key);
     return this.shardIds[index] || `shard_${index + 1}`;
   }
   getRandomShardIndex() {
-    return Math.floor(Math.random() * this.shardIds.length);
+    if (this.shardIds.length === 0) {
+      throw new RoutingError("No shards available");
+    }
+    let selection = Math.random() * this.totalRandomWeight;
+    for (let index = 0; index < this.randomShardWeights.length; index++) {
+      selection -= this.randomShardWeights[index];
+      if (selection < 0) {
+        return index;
+      }
+    }
+    return this.shardIds.length - 1;
   }
   getRandomShardId() {
     const index = this.getRandomShardIndex();
@@ -146,18 +175,235 @@ var ShardRouter = class {
   }
 };
 
+// src/constants/internal.ts
+var INTERNAL_DEFAULTS = {
+  HEALTH_CHECK_TIMEOUT_MS: 5e3,
+  CROSS_SHARD_CONCURRENCY: 4,
+  CROSS_SHARD_TIMEOUT_MS: 1e4,
+  CLI_COMMAND_TIMEOUT_MS: 12e4,
+  CLI_FORCE_KILL_GRACE_MS: 5e3,
+  CLI_MAX_OUTPUT_LENGTH: 1e6,
+  CLI_TEST_TCP_TIMEOUT_MS: 5e3,
+  CLI_TEST_COMMAND_TIMEOUT_MS: 15e3,
+  STUDIO_BASE_PORT: 51212,
+  STUDIO_REUSE_EXISTING: true,
+  STUDIO_STRICT_PORT_CHECK: false,
+  STUDIO_START_TIMEOUT_MS: 15e3,
+  STUDIO_STABILITY_MS: 500,
+  STUDIO_SHUTDOWN_TIMEOUT_MS: 5e3
+};
+
+// src/utils/sanitize.ts
+var escapeRegExp = (value) => value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+var maskDatabaseUrl = (url) => {
+  try {
+    const parsed = new URL(url);
+    if (parsed.password) {
+      parsed.password = "***";
+    }
+    return parsed.toString();
+  } catch {
+    return url.replace(/:[^:@]+@/, ":***@");
+  }
+};
+var sanitizeDatabaseText = (output, databaseUrls) => {
+  let sanitized = output;
+  for (const url of databaseUrls) {
+    sanitized = sanitized.replace(new RegExp(escapeRegExp(url), "g"), maskDatabaseUrl(url));
+    try {
+      const password = new URL(url).password;
+      if (password.length >= 4) {
+        sanitized = sanitized.replace(
+          new RegExp(escapeRegExp(decodeURIComponent(password)), "g"),
+          "***"
+        );
+      }
+    } catch {
+    }
+  }
+  return sanitized.replace(
+    /(postgres(?:ql)?:\/\/[^:/\s]+:)[^@\s]+(@)/gi,
+    "$1***$2"
+  );
+};
+
+// src/core/executor.ts
+var timeoutError = (shardId, timeoutMs) => new Error(`Shard ${shardId} operation timed out after ${timeoutMs}ms`);
+var CrossShardExecutor = class {
+  constructor(config = {}) {
+    this.metrics = {
+      operationCount: 0,
+      timeoutCount: 0,
+      failuresByShard: /* @__PURE__ */ new Map()
+    };
+    this.concurrency = config.concurrency ?? INTERNAL_DEFAULTS.CROSS_SHARD_CONCURRENCY;
+    this.timeoutMs = config.timeoutMs ?? INTERNAL_DEFAULTS.CROSS_SHARD_TIMEOUT_MS;
+  }
+  orderedTargets(targets) {
+    return targets.map((target, index) => ({ target, index })).sort((left, right) => {
+      if (left.target.isHealthy !== right.target.isHealthy) {
+        return left.target.isHealthy ? -1 : 1;
+      }
+      const leftLatency = left.target.latencyMs >= 0 ? left.target.latencyMs : Number.POSITIVE_INFINITY;
+      const rightLatency = right.target.latencyMs >= 0 ? right.target.latencyMs : Number.POSITIVE_INFINITY;
+      return leftLatency - rightLatency || left.index - right.index;
+    }).map(({ target }) => target);
+  }
+  recordFailure(shardId, error) {
+    this.metrics.failuresByShard.set(
+      shardId,
+      (this.metrics.failuresByShard.get(shardId) ?? 0) + 1
+    );
+    if (error.message.includes("timed out")) {
+      this.metrics.timeoutCount++;
+    }
+  }
+  withTimeout(shardId, operation) {
+    return new Promise((resolve, reject) => {
+      let settled = false;
+      const finish = (callback) => {
+        if (settled) {
+          return;
+        }
+        settled = true;
+        clearTimeout(timer);
+        callback();
+      };
+      const timer = setTimeout(() => {
+        finish(() => reject(timeoutError(shardId, this.timeoutMs)));
+      }, this.timeoutMs);
+      Promise.resolve().then(operation).then(
+        (result) => finish(() => resolve(result)),
+        (error) => finish(() => reject(error))
+      );
+    });
+  }
+  async executeAll(targets, operation) {
+    this.metrics.operationCount++;
+    const originalIndexes = new Map(targets.map((target, index) => [target.shardId, index]));
+    const ordered = this.orderedTargets(targets);
+    const results = new Array(targets.length);
+    let nextIndex = 0;
+    const worker = async () => {
+      while (nextIndex < ordered.length) {
+        const target = ordered[nextIndex++];
+        const resultIndex = originalIndexes.get(target.shardId);
+        try {
+          const result = await this.withTimeout(
+            target.shardId,
+            () => operation(target.client, target.shardId)
+          );
+          results[resultIndex] = { shardId: target.shardId, result, error: void 0 };
+        } catch (error) {
+          const normalizedError = error instanceof Error ? error : new Error(String(error));
+          this.recordFailure(target.shardId, normalizedError);
+          results[resultIndex] = {
+            shardId: target.shardId,
+            result: null,
+            error: normalizedError
+          };
+        }
+      }
+    };
+    const workerCount = Math.min(this.concurrency, ordered.length);
+    await Promise.all(Array.from({ length: workerCount }, () => worker()));
+    return results;
+  }
+  findFirst(targets, operation) {
+    this.metrics.operationCount++;
+    const ordered = this.orderedTargets(targets);
+    return new Promise((resolve) => {
+      if (ordered.length === 0) {
+        resolve({ result: null, shardId: null });
+        return;
+      }
+      let nextIndex = 0;
+      let active = 0;
+      let completed = 0;
+      let found = false;
+      const launch = () => {
+        while (!found && active < this.concurrency && nextIndex < ordered.length) {
+          const target = ordered[nextIndex++];
+          active++;
+          this.withTimeout(target.shardId, () => operation(target.client)).then((result) => {
+            if (!found && result !== null && result !== void 0) {
+              found = true;
+              resolve({ result, shardId: target.shardId });
+            }
+          }).catch((error) => {
+            const normalizedError = error instanceof Error ? error : new Error(String(error));
+            this.recordFailure(target.shardId, normalizedError);
+          }).finally(() => {
+            active--;
+            completed++;
+            if (!found && completed === ordered.length) {
+              resolve({ result: null, shardId: null });
+              return;
+            }
+            launch();
+          });
+        }
+      };
+      launch();
+    });
+  }
+};
+
 // src/core/manager.ts
+var isObject = (value) => typeof value === "object" && value !== null;
+var isConnectable = (client) => isObject(client) && typeof client.$connect === "function";
+var isDisconnectable = (client) => isObject(client) && typeof client.$disconnect === "function";
+var isQueryable = (client) => isObject(client) && typeof client.$queryRaw === "function";
+var withTimeout = (operation, timeoutMs, timeoutMessage) => new Promise((resolve, reject) => {
+  let settled = false;
+  const timer = setTimeout(() => {
+    if (!settled) {
+      settled = true;
+      reject(new Error(timeoutMessage));
+    }
+  }, timeoutMs);
+  operation.then(
+    (result) => {
+      if (!settled) {
+        settled = true;
+        clearTimeout(timer);
+        resolve(result);
+      }
+    },
+    (error) => {
+      if (!settled) {
+        settled = true;
+        clearTimeout(timer);
+        reject(error);
+      }
+    }
+  );
+});
 var ShardManager = class {
   constructor(config) {
     this.instances = /* @__PURE__ */ new Map();
+    this.unhealthyAccessWarnings = /* @__PURE__ */ new Set();
     this.healthCheckInterval = null;
+    this.healthCheckInFlight = false;
+    this.healthCheckPromise = null;
     this.config = config;
+    this.executor = new CrossShardExecutor();
   }
   async initialize() {
     this.config.logger.info(`Initializing ${this.config.shards.length} shard(s)...`);
-    for (const shardConfig of this.config.shards) {
-      this.initializeShard(shardConfig);
+    try {
+      for (const shardConfig of this.config.shards) {
+        this.initializeShard(shardConfig);
+      }
+    } catch (error) {
+      await this.disconnectInstances();
+      this.instances.clear();
+      throw error;
     }
+    void this.performHealthChecks(true).catch((error) => {
+      const message = error instanceof Error ? error.message : String(error);
+      this.config.logger.error(`Initial shard health verification failed: ${message}`);
+    });
     this.startHealthChecks();
     this.config.logger.info("All shards initialized successfully");
   }
@@ -179,7 +425,8 @@ var ShardManager = class {
       });
       this.config.logger.info(`Shard ${shardConfig.id} initialized`);
     } catch (error) {
-      this.config.logger.error(`Failed to initialize shard ${shardConfig.id}: ${error}`);
+      const message = sanitizeDatabaseText(String(error), [shardConfig.url]);
+      this.config.logger.error(`Failed to initialize shard ${shardConfig.id}: ${message}`);
       throw new ConnectionError(`Failed to initialize shard ${shardConfig.id}`, shardConfig.id);
     }
   }
@@ -187,58 +434,108 @@ var ShardManager = class {
     if (this.healthCheckInterval) {
       clearInterval(this.healthCheckInterval);
     }
-    this.healthCheckInterval = setInterval(async () => {
-      await this.performHealthChecks();
+    this.healthCheckInterval = setInterval(() => {
+      if (!this.healthCheckInFlight) {
+        void this.performHealthChecks(false);
+      }
     }, this.config.healthCheckIntervalMs);
+    this.healthCheckInterval.unref?.();
   }
-  async performHealthChecks() {
-    const checks = Array.from(this.instances.values()).map(async (instance) => {
-      const startTime = Date.now();
-      try {
-        const client = instance.client;
-        if (typeof client.$queryRaw === "function") {
-          await client.$queryRaw`SELECT 1`;
+  async verifyShard(instance, warmConnection) {
+    const startTime = Date.now();
+    const previousHealth = instance.health;
+    try {
+      const verification = async () => {
+        if (warmConnection && isConnectable(instance.client)) {
+          await instance.client.$connect();
         }
-        const latencyMs = Date.now() - startTime;
-        instance.health = {
-          ...instance.health,
-          isHealthy: true,
-          latencyMs,
-          lastChecked: /* @__PURE__ */ new Date(),
-          consecutiveFailures: 0
-        };
-      } catch (error) {
-        const consecutiveFailures = instance.health.consecutiveFailures + 1;
-        const isHealthy = consecutiveFailures < this.config.circuitBreakerThreshold;
-        instance.health = {
-          ...instance.health,
-          isHealthy,
-          latencyMs: -1,
-          lastChecked: /* @__PURE__ */ new Date(),
-          errorCount: instance.health.errorCount + 1,
-          consecutiveFailures
-        };
-        if (!isHealthy) {
-          this.config.logger.error(
-            `Shard ${instance.config.id} marked unhealthy after ${consecutiveFailures} consecutive failures`
-          );
+        if (isQueryable(instance.client)) {
+          await instance.client.$queryRaw`SELECT 1`;
         }
+      };
+      await withTimeout(
+        verification(),
+        INTERNAL_DEFAULTS.HEALTH_CHECK_TIMEOUT_MS,
+        `Health check timed out for shard ${instance.config.id}`
+      );
+      instance.health = {
+        ...previousHealth,
+        isHealthy: true,
+        latencyMs: Date.now() - startTime,
+        lastChecked: /* @__PURE__ */ new Date(),
+        consecutiveFailures: 0
+      };
+      this.unhealthyAccessWarnings.delete(instance.config.id);
+      if (!previousHealth.isHealthy && previousHealth.consecutiveFailures > 0) {
+        this.config.logger.info(`Shard ${instance.config.id} recovered`);
       }
+    } catch (error) {
+      const consecutiveFailures = previousHealth.consecutiveFailures + 1;
+      const wasVerifiedHealthy = previousHealth.isHealthy;
+      const isHealthy = warmConnection ? false : wasVerifiedHealthy && consecutiveFailures < this.config.circuitBreakerThreshold;
+      instance.health = {
+        ...previousHealth,
+        isHealthy,
+        latencyMs: -1,
+        lastChecked: /* @__PURE__ */ new Date(),
+        errorCount: previousHealth.errorCount + 1,
+        consecutiveFailures
+      };
+      const message = sanitizeDatabaseText(
+        error instanceof Error ? error.message : String(error),
+        [instance.config.url]
+      );
+      const isFirstFailure = previousHealth.errorCount === 0;
+      const becameUnhealthy = previousHealth.isHealthy && !isHealthy;
+      if (isFirstFailure || becameUnhealthy) {
+        this.config.logger.error(
+          `Health check failed for ${instance.config.id}: ${message}${!isHealthy ? "; shard marked unhealthy" : ""}`
+        );
+      }
+    }
+  }
+  performHealthChecks(warmConnections) {
+    if (this.healthCheckInFlight) {
+      return this.healthCheckPromise ?? Promise.resolve();
+    }
+    this.healthCheckInFlight = true;
+    const operation = this.executor.executeAll(this.getExecutionTargets(), async (_client, shardId) => {
+      const instance = this.instances.get(shardId);
+      if (instance) {
+        await this.verifyShard(instance, warmConnections);
+      }
+      return null;
+    }).then(() => void 0).finally(() => {
+      this.healthCheckInFlight = false;
+      this.healthCheckPromise = null;
     });
-    await Promise.allSettled(checks);
+    this.healthCheckPromise = operation;
+    return operation;
+  }
+  getExecutionTargets() {
+    return Array.from(this.instances.values()).map((instance) => ({
+      shardId: instance.config.id,
+      client: instance.client,
+      isHealthy: instance.health.isHealthy,
+      latencyMs: instance.health.latencyMs
+    }));
   }
   getClient(shardId) {
     const instance = this.instances.get(shardId);
     if (!instance) {
       throw new ConnectionError(`Shard ${shardId} not found`, shardId);
     }
-    if (!instance.health.isHealthy) {
+    if (!instance.health.isHealthy && !this.unhealthyAccessWarnings.has(shardId)) {
+      this.unhealthyAccessWarnings.add(shardId);
       this.config.logger.warn(`Accessing unhealthy shard ${shardId}`);
     }
     return instance.client;
   }
   getClientByIndex(index) {
-    const shardId = `shard_${index + 1}`;
+    const shardId = this.getShardIds()[index];
+    if (!shardId) {
+      throw new ConnectionError(`Shard index ${index} not found`, String(index));
+    }
     return this.getClient(shardId);
   }
   getAllClients() {
@@ -259,56 +556,38 @@ var ShardManager = class {
   getAllHealth() {
     return Array.from(this.instances.values()).map((instance) => instance.health);
   }
-  async executeOnAll(operation) {
-    const results = await Promise.allSettled(
-      Array.from(this.instances.entries()).map(async ([shardId, instance]) => {
-        try {
-          const result = await operation(instance.client, shardId);
-          return { shardId, result, error: void 0 };
-        } catch (error) {
-          return { shardId, result: null, error };
-        }
-      })
-    );
-    return results.map((result) => {
-      if (result.status === "fulfilled") {
-        return result.value;
-      }
-      return {
-        shardId: "unknown",
-        result: null,
-        error: result.reason
-      };
-    });
+  executeOnAll(operation) {
+    return this.executor.executeAll(this.getExecutionTargets(), operation);
   }
-  async findFirst(operation) {
-    const results = await this.executeOnAll(operation);
-    for (const res of results) {
-      if (res.result !== null && res.result !== void 0) {
-        return { result: res.result, shardId: res.shardId };
-      }
-    }
-    return { result: null, shardId: null };
+  findFirst(operation) {
+    return this.executor.findFirst(this.getExecutionTargets(), operation);
   }
-  async shutdown() {
-    this.config.logger.info("Shutting down...");
-    if (this.healthCheckInterval) {
-      clearInterval(this.healthCheckInterval);
-      this.healthCheckInterval = null;
-    }
+  async disconnectInstances() {
     const disconnects = Array.from(this.instances.values()).map(async (instance) => {
       try {
-        const client = instance.client;
-        if (typeof client.$disconnect === "function") {
-          await client.$disconnect();
+        if (isDisconnectable(instance.client)) {
+          await instance.client.$disconnect();
         }
         this.config.logger.info(`Shard ${instance.config.id} disconnected`);
       } catch (error) {
-        this.config.logger.error(`Error disconnecting shard ${instance.config.id}: ${error}`);
+        const message = sanitizeDatabaseText(String(error), [instance.config.url]);
+        this.config.logger.error(`Error disconnecting shard ${instance.config.id}: ${message}`);
       }
     });
     await Promise.allSettled(disconnects);
+  }
+  async shutdown() {
+    this.config.logger.info("Shutting down...");
+    if (this.healthCheckInterval) {
+      clearInterval(this.healthCheckInterval);
+      this.healthCheckInterval = null;
+    }
+    if (this.healthCheckPromise) {
+      await this.healthCheckPromise;
+    }
+    await this.disconnectInstances();
     this.instances.clear();
+    this.unhealthyAccessWarnings.clear();
     this.config.logger.info("Shutdown complete");
   }
 };
@@ -327,17 +606,34 @@ var PrismaSharding = class {
     if (!config.shards || config.shards.length === 0) {
       throw new ConfigError(ERROR_MESSAGES.NO_SHARDS);
     }
-    if (!config.createClient) {
+    if (typeof config.createClient !== "function") {
       throw new ConfigError(ERROR_MESSAGES.MISSING_CLIENT_FACTORY);
     }
+    const shardIds = /* @__PURE__ */ new Set();
     for (const shard of config.shards) {
+      if (!shard.id || shard.id.trim().length === 0) {
+        throw new ConfigError("Shard ID must not be empty");
+      }
+      if (shardIds.has(shard.id)) {
+        throw new ConfigError(`Duplicate shard ID: "${shard.id}"`);
+      }
+      shardIds.add(shard.id);
       if (!shard.url || !validateUrl(shard.url)) {
         throw new ConfigError(ERROR_MESSAGES.INVALID_SHARD_URL(shard.id));
       }
+      if (shard.weight !== void 0 && (!Number.isFinite(shard.weight) || shard.weight <= 0)) {
+        throw new ConfigError(`Weight for shard "${shard.id}" must be positive`);
+      }
     }
     if (config.strategy && config.strategy !== "modulo" && config.strategy !== "consistent-hash") {
       throw new ConfigError(ERROR_MESSAGES.INVALID_STRATEGY(config.strategy));
     }
+    if (config.healthCheckIntervalMs !== void 0 && (!Number.isFinite(config.healthCheckIntervalMs) || config.healthCheckIntervalMs <= 0)) {
+      throw new ConfigError("healthCheckIntervalMs must be positive");
+    }
+    if (config.circuitBreakerThreshold !== void 0 && (!Number.isFinite(config.circuitBreakerThreshold) || config.circuitBreakerThreshold <= 0)) {
+      throw new ConfigError("circuitBreakerThreshold must be positive");
+    }
   }
   async connect() {
     if (this.connected) {
@@ -351,10 +647,20 @@ var PrismaSharding = class {
       circuitBreakerThreshold: this.config.circuitBreakerThreshold ?? DEFAULTS.CIRCUIT_BREAKER_THRESHOLD,
       logger: this.logger
     });
-    await this.manager.initialize();
+    try {
+      await this.manager.initialize();
+    } catch (error) {
+      this.manager = null;
+      this.router = null;
+      this.connected = false;
+      throw error;
+    }
     this.router = new ShardRouter({
       strategy: this.config.strategy ?? "modulo",
       shardIds: this.manager.getShardIds(),
+      shardWeights: new Map(
+        this.config.shards.map((shard) => [shard.id, shard.weight ?? 1])
+      ),
       logger: this.logger
     });
     this.connected = true;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "prisma-sharding",
-  "version": "0.0.6",
+  "version": "0.0.7",
   "description": "Lightweight database sharding library for Prisma with connection pooling and health monitoring",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -35,9 +35,9 @@
     }
   },
   "scripts": {
-    "build": "tsup src/index.ts --format cjs,esm --dts --clean && tsup src/cli/migrate.ts src/cli/studio.ts src/cli/test.ts src/cli/update.ts --format cjs --outDir dist/cli --clean",
+    "build": "tsup src/index.ts --format cjs,esm --dts --clean && tsup src/cli/migrate.ts src/cli/studio.ts src/cli/test.ts src/cli/update.ts src/cli/utils/command.ts src/cli/utils/postgres.ts --format cjs --outDir dist/cli --clean",
     "build:lib": "tsup src/index.ts --format cjs,esm --dts --clean",
-    "build:cli": "tsup src/cli/migrate.ts src/cli/studio.ts src/cli/test.ts src/cli/update.ts --format cjs --outDir dist/cli",
+    "build:cli": "tsup src/cli/migrate.ts src/cli/studio.ts src/cli/test.ts src/cli/update.ts src/cli/utils/command.ts src/cli/utils/postgres.ts --format cjs --outDir dist/cli",
     "dev": "tsup src/index.ts --format cjs,esm --dts --watch",
     "lint": "eslint src --ext .ts",
     "test": "yarn build && node --test test/*.test.js",