Skip to content

feat: add atlas-connect-cluster tool #131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
Apr 28, 2025
Merged
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@ You may experiment asking `Can you connect to my mongodb instance?`.
- `atlas-list-clusters` - Lists MongoDB Atlas clusters
- `atlas-inspect-cluster` - Inspect a specific MongoDB Atlas cluster
- `atlas-create-free-cluster` - Create a free MongoDB Atlas cluster
- `atlas-connect-cluster` - Connects to MongoDB Atlas cluster
- `atlas-inspect-access-list` - Inspect IP/CIDR ranges with access to MongoDB Atlas clusters
- `atlas-create-access-list` - Configure IP/CIDR access list for MongoDB Atlas clusters
- `atlas-list-db-users` - List MongoDB Atlas database users
Expand Down
2 changes: 2 additions & 0 deletions src/logger.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ export const LogId = {
serverInitialized: mongoLogId(1_000_002),

atlasCheckCredentials: mongoLogId(1_001_001),
atlasDeleteDatabaseUserFailure: mongoLogId(1_001_002),

telemetryDisabled: mongoLogId(1_002_001),
telemetryEmitFailure: mongoLogId(1_002_002),
Expand All @@ -22,6 +23,7 @@ export const LogId = {
toolDisabled: mongoLogId(1_003_003),

mongodbConnectFailure: mongoLogId(1_004_001),
mongodbDisconnectFailure: mongoLogId(1_004_002),
} as const;

abstract class LoggerBase {
Expand Down
46 changes: 42 additions & 4 deletions src/session.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
import { ApiClient, ApiClientCredentials } from "./common/atlas/apiClient.js";
import { Implementation } from "@modelcontextprotocol/sdk/types.js";
import logger, { LogId } from "./logger.js";
import EventEmitter from "events";
import { ConnectOptions } from "./config.js";

Expand All @@ -12,6 +13,7 @@ export interface SessionOptions {

export class Session extends EventEmitter<{
close: [];
disconnect: [];
}> {
sessionId?: string;
serviceProvider?: NodeDriverServiceProvider;
Expand All @@ -20,6 +22,12 @@ export class Session extends EventEmitter<{
name: string;
version: string;
};
connectedAtlasCluster?: {
username: string;
projectId: string;
clusterName: string;
expiryDate: Date;
};

constructor({ apiBaseUrl, apiClientId, apiClientSecret }: SessionOptions) {
super();
Expand Down Expand Up @@ -47,17 +55,47 @@ export class Session extends EventEmitter<{
}
}

async close(): Promise<void> {
async disconnect(): Promise<void> {
if (this.serviceProvider) {
try {
await this.serviceProvider.close(true);
} catch (error) {
console.error("Error closing service provider:", error);
} catch (err: unknown) {
const error = err instanceof Error ? err : new Error(String(err));
logger.error(LogId.mongodbDisconnectFailure, "Error closing service provider:", error.message);
}
this.serviceProvider = undefined;
}
if (!this.connectedAtlasCluster) {
this.emit("disconnect");
return;
}
try {
await this.apiClient.deleteDatabaseUser({
params: {
path: {
groupId: this.connectedAtlasCluster.projectId,
username: this.connectedAtlasCluster.username,
databaseName: "admin",
},
},
});
} catch (err: unknown) {
const error = err instanceof Error ? err : new Error(String(err));

this.emit("close");
logger.error(
LogId.atlasDeleteDatabaseUserFailure,
"atlas-connect-cluster",
`Error deleting previous database user: ${error.message}`
);
}
this.connectedAtlasCluster = undefined;

this.emit("disconnect");
}

async close(): Promise<void> {
await this.disconnect();
this.emit("close");
}

async connectToMongoDB(connectionString: string, connectOptions: ConnectOptions): Promise<void> {
Expand Down
5 changes: 4 additions & 1 deletion src/tools/atlas/create/createFreeCluster.ts
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,10 @@ export class CreateFreeClusterTool extends AtlasToolBase {
});

return {
content: [{ type: "text", text: `Cluster "${name}" has been created in region "${region}".` }],
content: [
{ type: "text", text: `Cluster "${name}" has been created in region "${region}".` },
{ type: "text", text: `Double check your access lists to enable your current IP.` },
],
};
}
}
114 changes: 114 additions & 0 deletions src/tools/atlas/metadata/connectCluster.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
import { z } from "zod";
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { AtlasToolBase } from "../atlasTool.js";
import { ToolArgs, OperationType } from "../../tool.js";
import { randomBytes } from "crypto";
import { promisify } from "util";

const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours

const randomBytesAsync = promisify(randomBytes);

async function generateSecurePassword(): Promise<string> {
const buf = await randomBytesAsync(16);
const pass = buf.toString("base64url");
return pass;
}

export class ConnectClusterTool extends AtlasToolBase {
protected name = "atlas-connect-cluster";
protected description = "Connect to MongoDB Atlas cluster";
protected operationType: OperationType = "metadata";
protected argsShape = {
projectId: z.string().describe("Atlas project ID"),
clusterName: z.string().describe("Atlas cluster name"),
};

protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
await this.session.disconnect();

const cluster = await this.session.apiClient.getCluster({
params: {
path: {
groupId: projectId,
clusterName,
},
},
});

if (!cluster) {
throw new Error("Cluster not found");
}

const baseConnectionString = cluster.connectionStrings?.standardSrv || cluster.connectionStrings?.standard;

if (!baseConnectionString) {
throw new Error("Connection string not available");
}

const username = `mcpUser${Math.floor(Math.random() * 100000)}`;
const password = await generateSecurePassword();

const expiryDate = new Date(Date.now() + EXPIRY_MS);

const readOnly =
this.config.readOnly ||
(this.config.disabledTools?.includes("create") &&
this.config.disabledTools?.includes("update") &&
this.config.disabledTools?.includes("delete") &&
!this.config.disabledTools?.includes("read") &&
!this.config.disabledTools?.includes("metadata"));

const roleName = readOnly ? "readAnyDatabase" : "readWriteAnyDatabase";

await this.session.apiClient.createDatabaseUser({
params: {
path: {
groupId: projectId,
},
},
body: {
databaseName: "admin",
groupId: projectId,
roles: [
{
roleName,
databaseName: "admin",
},
],
scopes: [{ type: "CLUSTER", name: clusterName }],
username,
password,
awsIAMType: "NONE",
ldapAuthType: "NONE",
oidcAuthType: "NONE",
x509Type: "NONE",
deleteAfterDate: expiryDate.toISOString(),
},
});

this.session.connectedAtlasCluster = {
username,
projectId,
clusterName,
expiryDate,
};

const cn = new URL(baseConnectionString);
cn.username = username;
cn.password = password;
cn.searchParams.set("authSource", "admin");
const connectionString = cn.toString();

await this.session.connectToMongoDB(connectionString, this.config.connectOptions);

return {
content: [
{
type: "text",
text: `Connected to cluster "${clusterName}"`,
},
],
};
}
}
2 changes: 2 additions & 0 deletions src/tools/atlas/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import { ListDBUsersTool } from "./read/listDBUsers.js";
import { CreateDBUserTool } from "./create/createDBUser.js";
import { CreateProjectTool } from "./create/createProject.js";
import { ListOrganizationsTool } from "./read/listOrgs.js";
import { ConnectClusterTool } from "./metadata/connectCluster.js";

export const AtlasTools = [
ListClustersTool,
Expand All @@ -20,4 +21,5 @@ export const AtlasTools = [
CreateDBUserTool,
CreateProjectTool,
ListOrganizationsTool,
ConnectClusterTool,
];
4 changes: 0 additions & 4 deletions tests/integration/tools/atlas/atlasHelpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,6 @@ import { config } from "../../../../src/config.js";

export type IntegrationTestFunction = (integration: IntegrationTest) => void;

export function sleep(ms: number) {
return new Promise((resolve) => setTimeout(resolve, ms));
}

export function describeWithAtlas(name: string, fn: IntegrationTestFunction) {
const testDefinition = () => {
const integration = setupIntegrationTest(() => ({
Expand Down
72 changes: 68 additions & 4 deletions tests/integration/tools/atlas/clusters.test.ts
Original file line number Diff line number Diff line change
@@ -1,14 +1,18 @@
import { Session } from "../../../../src/session.js";
import { expectDefined } from "../../helpers.js";
import { describeWithAtlas, withProject, sleep, randomId } from "./atlasHelpers.js";
import { describeWithAtlas, withProject, randomId } from "./atlasHelpers.js";
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";

function sleep(ms: number) {
return new Promise((resolve) => setTimeout(resolve, ms));
}

async function deleteAndWaitCluster(session: Session, projectId: string, clusterName: string) {
await session.apiClient.deleteCluster({
params: {
path: {
groupId: projectId,
clusterName: clusterName,
clusterName,
},
},
});
Expand All @@ -18,7 +22,7 @@ async function deleteAndWaitCluster(session: Session, projectId: string, cluster
params: {
path: {
groupId: projectId,
clusterName: clusterName,
clusterName,
},
},
});
Expand All @@ -29,6 +33,23 @@ async function deleteAndWaitCluster(session: Session, projectId: string, cluster
}
}

async function waitClusterState(session: Session, projectId: string, clusterName: string, state: string) {
while (true) {
const cluster = await session.apiClient.getCluster({
params: {
path: {
groupId: projectId,
clusterName,
},
},
});
if (cluster?.stateName === state) {
return;
}
await sleep(1000);
}
}

describeWithAtlas("clusters", (integration) => {
withProject(integration, ({ getProjectId }) => {
const clusterName = "ClusterTest-" + randomId;
Expand Down Expand Up @@ -66,7 +87,7 @@ describeWithAtlas("clusters", (integration) => {
},
})) as CallToolResult;
expect(response.content).toBeArray();
expect(response.content).toHaveLength(1);
expect(response.content).toHaveLength(2);
expect(response.content[0].text).toContain("has been created");
});
});
Expand Down Expand Up @@ -117,5 +138,48 @@ describeWithAtlas("clusters", (integration) => {
expect(response.content[1].text).toContain(`${clusterName} | `);
});
});

describe("atlas-connect-cluster", () => {
beforeAll(async () => {
const projectId = getProjectId();
await waitClusterState(integration.mcpServer().session, projectId, clusterName, "IDLE");
await integration.mcpServer().session.apiClient.createProjectIpAccessList({
params: {
path: {
groupId: projectId,
},
},
body: [
{
comment: "MCP test",
cidrBlock: "0.0.0.0/0",
},
],
});
});

it("should have correct metadata", async () => {
const { tools } = await integration.mcpClient().listTools();
const connectCluster = tools.find((tool) => tool.name === "atlas-connect-cluster");

expectDefined(connectCluster);
expect(connectCluster.inputSchema.type).toBe("object");
expectDefined(connectCluster.inputSchema.properties);
expect(connectCluster.inputSchema.properties).toHaveProperty("projectId");
expect(connectCluster.inputSchema.properties).toHaveProperty("clusterName");
});

it("connects to cluster", async () => {
const projectId = getProjectId();

const response = (await integration.mcpClient().callTool({
name: "atlas-connect-cluster",
arguments: { projectId, clusterName },
})) as CallToolResult;
expect(response.content).toBeArray();
expect(response.content).toHaveLength(1);
expect(response.content[0].text).toContain(`Connected to cluster "${clusterName}"`);
});
});
});
});
Loading