mongoku
Version:
[](https://github.com/huggingface/Mongoku/actions/workflows/ci.yml)
392 lines (389 loc) • 13.1 kB
JavaScript
import { p as private_env } from './shared-server-BmU87nph.js';
import { l as logger } from './logger-PfH_grbh.js';
import { resolveSrv } from 'dns/promises';
import { ReadPreference, MongoClient } from 'mongodb';
import { URL } from 'url';
import * as fs from 'node:fs';
import * as os from 'node:os';
import * as path from 'node:path';
const DEFAULT_HOSTS = private_env.MONGOKU_DEFAULT_HOST ? private_env.MONGOKU_DEFAULT_HOST.split(";") : ["localhost:27017"];
const DATABASE_FILE = private_env.MONGOKU_DATABASE_FILE || path.join(os.homedir(), ".mongoku.db");
class HostsManager {
_hosts = /* @__PURE__ */ new Map();
// path -> _id
async load() {
let first = false;
try {
await fs.promises.stat(DATABASE_FILE);
} catch {
first = true;
}
if (!first) {
await this._loadFromFile();
}
if (first || this._hosts.size === 0) {
for (const hostname of DEFAULT_HOSTS) {
this._hosts.set(hostname, this._generateId());
}
await this._saveToFile();
}
}
async _loadFromFile() {
const content = await fs.promises.readFile(DATABASE_FILE, "utf8");
const lines = content.trim().split("\n").filter((line) => line.trim());
const newHosts = /* @__PURE__ */ new Map();
for (const line of lines) {
const host = JSON.parse(line);
if (host && typeof host.path === "string") {
const id = host._id || this._generateId();
newHosts.set(host.path, id);
}
}
this._hosts = newHosts;
}
async _saveToFile() {
const lines = Array.from(this._hosts).map(([hostPath, id]) => JSON.stringify({ path: hostPath, _id: id }));
await fs.promises.writeFile(DATABASE_FILE, lines.join("\n") + "\n", "utf8");
}
_generateId() {
const chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
let result = "";
for (let i = 0; i < 16; i++) {
result += chars.charAt(Math.floor(Math.random() * chars.length));
}
return result;
}
async getHosts() {
return Array.from(this._hosts).map(([hostPath, id]) => ({ path: hostPath, _id: id }));
}
getHost(name) {
return this._hosts.get(name);
}
async add(hostPath) {
let id = this._hosts.get(hostPath);
if (!id) {
id = this._generateId();
this._hosts.set(hostPath, id);
}
await this._saveToFile();
return id;
}
async removeById(id) {
for (const [hostPath, hostId] of this._hosts) {
if (hostId === id) {
this._hosts.delete(hostPath);
break;
}
}
await this._saveToFile();
}
}
async function getCollectionJson(collection, type) {
const stats = {
size: 0,
count: 0,
avgObjSize: 0,
storageSize: 0,
capped: false,
nindexes: 0,
totalIndexSize: 0,
indexSizes: {}
};
if (type !== "view") {
const agg = await collection.aggregate([
{
$collStats: {
storageStats: {},
count: {}
}
}
]).next().catch(() => null);
if (agg) {
stats.size = agg.storageStats.size;
stats.count = agg.storageStats.count;
stats.avgObjSize = Math.round(agg.storageStats.size / agg.storageStats.count);
stats.storageSize = agg.storageStats.storageSize;
stats.capped = agg.storageStats.capped;
stats.nindexes = agg.storageStats.nindexes;
stats.totalIndexSize = agg.storageStats.totalIndexSize;
stats.indexSizes = agg.storageStats.indexSizes;
}
}
let indexes = [];
if (type !== "view") {
try {
const indexList = await collection.listIndexes().toArray();
indexes = indexList.map((index) => ({
name: index.name,
key: index.key,
size: stats.indexSizes[index.name] || 0
}));
} catch {
indexes = Object.entries(stats.indexSizes).map(([name, size]) => ({
name,
size
}));
}
}
return {
name: collection.collectionName,
size: (stats.storageSize ?? 0) + (stats.totalIndexSize ?? 0),
dataSize: stats.size,
count: stats.count,
avgObjSize: stats.avgObjSize ?? 0,
storageSize: stats.storageSize ?? 0,
capped: stats.capped,
nIndexes: stats.nindexes,
totalIndexSize: stats.totalIndexSize ?? 0,
indexSizes: stats.indexSizes,
indexes
};
}
async function extractNodesFromConnectionString(connectionString) {
const url = new URL(connectionString);
if (url.protocol === "mongodb+srv:") {
const hostname = url.hostname;
const srvRecords = await resolveSrv(`_mongodb._tcp.${hostname}`);
return srvRecords.sort((a, b) => a.priority - b.priority || b.weight - a.weight).map((record) => `${record.name}:${record.port}`).sort();
}
if (url.protocol === "mongodb:") {
const hostsString = url.host;
if (!hostsString) {
throw new Error("No hosts found in connection string");
}
return hostsString.split(",").map((host) => host.trim()).filter((host) => host.length > 0).sort();
}
throw new TypeError(`Unsupported protocol: ${url.protocol}`);
}
class MongoClientWithMappings extends MongoClient {
url;
mappings = {};
indexes = {};
_id;
name;
constructor(url, _id, name, readPreference) {
super(url, readPreference ? { readPreference } : {});
this.url = url;
this._id = _id;
this.name = name;
}
async getMappings(dbName, collectionName, opts) {
if (opts?.forceRefresh || !this.mappings[dbName]?.[collectionName]) {
const mappings = await this.db(dbName).collection("mongoku.mappings").findOne({ _id: collectionName });
this.mappings[dbName] ??= {};
this.mappings[dbName][collectionName] = mappings?.mappings ?? {};
}
return this.mappings[dbName][collectionName];
}
clearMappingsCache(dbName, collectionName) {
delete this.mappings[dbName][collectionName];
}
async getIndexes(dbName, collectionName, opts) {
if (opts?.forceRefresh || !this.indexes[dbName]?.[collectionName]) {
const indexList = await this.db(dbName).collection(collectionName).listIndexes().toArray();
this.indexes[dbName] ??= {};
this.indexes[dbName][collectionName] = indexList.map((index) => ({
name: index.name,
key: index.key
}));
}
return this.indexes[dbName][collectionName];
}
setIndexes(dbName, collectionName, indexes) {
this.indexes[dbName] ??= {};
this.indexes[dbName][collectionName] = indexes;
}
clearIndexesCache(dbName, collectionName) {
delete this.indexes[dbName]?.[collectionName];
}
/**
* Check if a field has an index (as the first key in the index)
*/
async hasIndexOnField(dbName, collectionName, field) {
const indexes = await this.getIndexes(dbName, collectionName);
return indexes.some((index) => {
const keys = Object.keys(index.key);
return keys[0] === field;
});
}
}
class MongoConnections {
/**
* Todo: better system where we can have mutiple servers with same hostname, and labels for each server that
* would be displayed in the UI instead of the hostname.
*/
clients = /* @__PURE__ */ new Map();
// _id -> MongoClientWithMappings
clientIds = /* @__PURE__ */ new Map();
// hostname -> _id
hostsManager;
countTimeout = parseInt(private_env.MONGOKU_COUNT_TIMEOUT, 10) || 3e4;
queryTimeout = private_env.MONGOKU_QUERY_TIMEOUT ? parseInt(private_env.MONGOKU_QUERY_TIMEOUT, 10) : void 0;
excludedDatabases;
readPreference;
constructor() {
this.hostsManager = new HostsManager();
const excludeEnv = private_env.MONGOKU_EXCLUDE_DATABASES || "";
this.excludedDatabases = new Set(
excludeEnv.split(",").map((db) => db.trim()).filter((db) => db.length > 0)
);
const readPrefMode = private_env.MONGOKU_READ_PREFERENCE;
if (readPrefMode) {
let tags;
if (private_env.MONGOKU_READ_PREFERENCE_TAGS) {
try {
tags = JSON.parse(private_env.MONGOKU_READ_PREFERENCE_TAGS);
} catch (err) {
logger.error("Failed to parse MONGOKU_READ_PREFERENCE_TAGS:", err);
}
}
this.readPreference = tags ? new ReadPreference(readPrefMode, tags) : new ReadPreference(readPrefMode);
logger.log(`Read preference configured: ${readPrefMode}${tags ? ` with tags ${JSON.stringify(tags)}` : ""}`);
}
}
async initialize() {
await this.hostsManager.load();
const hosts = await this.hostsManager.getHosts();
for (const host of hosts) {
const urlStr = host.path.startsWith("mongodb") ? host.path : `mongodb://${host.path}`;
try {
const url = new URL(urlStr);
const hostname = url.host || host.path;
if (!this.clients.has(hostname)) {
const client = new MongoClientWithMappings(urlStr, host._id, hostname, this.readPreference);
this.clients.set(host._id, client);
this.clientIds.set(hostname, host._id);
}
} catch (err) {
logger.error(`Failed to parse URL for host ${host.path}:`, err);
}
}
}
getClient(name) {
const clientId = this.clientIds.get(name) || this.clientIds.get(`${name}:27017`);
if (!clientId) {
throw new Error(`Client not found: ${name}`);
}
const client = this.clients.get(clientId);
if (!client) {
throw new Error(`Client not found: ${name}`);
}
return client;
}
listClients() {
return Array.from(this.clients.values()).map((client) => ({
name: client.name,
_id: client._id,
client
}));
}
getCountTimeout() {
return this.countTimeout;
}
getQueryTimeout() {
return this.queryTimeout;
}
filterDatabases(databases) {
if (this.excludedDatabases.size === 0) {
return databases;
}
return databases.filter((db) => !this.excludedDatabases.has(db.name));
}
async addServer(hostPath) {
const id = await this.hostsManager.add(hostPath);
const urlStr = hostPath.startsWith("mongodb") ? hostPath : `mongodb://${hostPath}`;
try {
const url = new URL(urlStr);
const hostname = url.host || hostPath;
if (!this.clients.has(hostname)) {
const client = new MongoClientWithMappings(urlStr, id, hostname, this.readPreference);
this.clients.set(id, client);
this.clientIds.set(hostname, id);
}
} catch (err) {
logger.error(`Failed to parse URL for host ${hostPath}:`, err);
throw err;
}
}
async removeServer(name) {
const clientKey = this.clientIds.has(name) ? name : `${name}:27017`;
const clientId = this.clientIds.get(clientKey);
if (!clientId) {
throw new Error(`Server not found: ${name}`);
}
await this.hostsManager.removeById(clientId);
this.clients.get(clientId)?.close().catch((err) => logger.error(`Error closing client ${name}:`, err));
this.clients.delete(clientId);
this.clientIds.delete(clientKey);
}
async reconnectClient(id) {
const oldClient = this.clients.get(id);
if (!oldClient) {
throw new Error(`Client not found: ${id}`);
}
oldClient.close().catch((err) => logger.error(`Error closing old client ${id}:`, err));
const newClient = new MongoClientWithMappings(oldClient.url, id, oldClient.name, this.readPreference);
this.clients.set(id, newClient);
await newClient.connect();
}
/**
* Get the list of nodes from a server's connection string
*/
async getServerNodes(serverId) {
const client = this.getClient(serverId);
return extractNodesFromConnectionString(client.url);
}
/**
* Fetch index stats from a specific node using directConnection
*/
async getIndexStatsFromNode(serverId, node, database, collection) {
const client = this.getClient(serverId);
const url = new URL(client.url);
const directUrl = new URL(url.toString());
directUrl.protocol = "mongodb:";
directUrl.host = node;
directUrl.searchParams.set("directConnection", "true");
if (url.searchParams.has("tls") || url.searchParams.has("ssl") || url.protocol === "mongodb+srv:") {
directUrl.searchParams.set("tls", "true");
}
let nodeClient = null;
try {
nodeClient = new MongoClient(directUrl.toString());
await nodeClient.connect();
const coll = nodeClient.db(database).collection(collection);
const statsResult = await coll.aggregate([{ $indexStats: {} }]).toArray();
const indexStats = Object.fromEntries(
statsResult.map((stat) => [
stat.name,
{
ops: stat.accesses?.ops || 0,
since: stat.accesses?.since || /* @__PURE__ */ new Date(),
host: stat.host || node
}
])
);
return indexStats;
} finally {
if (nodeClient) {
await nodeClient.close();
}
}
}
}
let mongoConnections = null;
let initPromise = null;
async function getMongo() {
if (!mongoConnections) {
if (!initPromise) {
initPromise = (async () => {
mongoConnections = new MongoConnections();
await mongoConnections.initialize();
return mongoConnections;
})();
}
return initPromise;
}
return mongoConnections;
}
export { getCollectionJson as a, getMongo as g };
//# sourceMappingURL=mongo-B92d7zNj.js.map