Skip to content

Commit 5dd6948

Browse files
himanshusinghshimanshusinghs
committed
chore: extract abstract ConnectionManager
This commit extracts an abstract class ConnectionManager out of MCPConnectionManager and modifies the TransportRunner interface to have the ConnectionManager implementation injected through a factory function. Contains a small drive by fix for making the ConnectionManager event emitting internal to the class itself.
1 parent 35b4a12 commit 5dd6948

File tree

15 files changed

+195
-157
lines changed

15 files changed

+195
-157
lines changed

‎package-lock.json

Lines changed: 0 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/common/connectionManager.ts

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
import { EventEmitter } from "events";
2+
import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
3+
4+
export interface AtlasClusterConnectionInfo {
5+
username: string;
6+
projectId: string;
7+
clusterName: string;
8+
expiryDate: Date;
9+
}
10+
11+
type ConnectionTag = "connected" | "connecting" | "disconnected" | "errored";
12+
export type OIDCConnectionAuthType = "oidc-auth-flow" | "oidc-device-flow";
13+
export type ConnectionStringAuthType = "scram" | "ldap" | "kerberos" | OIDCConnectionAuthType | "x.509";
14+
15+
export interface ConnectionState {
16+
tag: ConnectionTag;
17+
connectionStringAuthType?: ConnectionStringAuthType;
18+
connectedAtlasCluster?: AtlasClusterConnectionInfo;
19+
}
20+
21+
export interface ConnectionStateConnected extends ConnectionState {
22+
tag: "connected";
23+
serviceProvider: NodeDriverServiceProvider;
24+
}
25+
26+
export interface ConnectionStateConnecting extends ConnectionState {
27+
tag: "connecting";
28+
serviceProvider: NodeDriverServiceProvider;
29+
oidcConnectionType: OIDCConnectionAuthType;
30+
oidcLoginUrl?: string;
31+
oidcUserCode?: string;
32+
}
33+
34+
export interface ConnectionStateDisconnected extends ConnectionState {
35+
tag: "disconnected";
36+
}
37+
38+
export interface ConnectionStateErrored extends ConnectionState {
39+
tag: "errored";
40+
errorReason: string;
41+
}
42+
43+
export type AnyConnectionState =
44+
| ConnectionStateConnected
45+
| ConnectionStateConnecting
46+
| ConnectionStateDisconnected
47+
| ConnectionStateErrored;
48+
49+
export interface ConnectionManagerEvents {
50+
"connection-requested": [AnyConnectionState];
51+
"connection-succeeded": [ConnectionStateConnected];
52+
"connection-timed-out": [ConnectionStateErrored];
53+
"connection-closed": [ConnectionStateDisconnected];
54+
"connection-errored": [ConnectionStateErrored];
55+
}
56+
57+
export interface MCPConnectParams {
58+
connectionString: string;
59+
atlas?: AtlasClusterConnectionInfo;
60+
}
61+
62+
export abstract class ConnectionManager<ConnectParams extends MCPConnectParams = MCPConnectParams> {
63+
protected clientName: string = "unknown";
64+
65+
protected readonly _events = new EventEmitter<ConnectionManagerEvents>();
66+
readonly events: Pick<EventEmitter<ConnectionManagerEvents>, "on" | "off" | "once"> = this._events;
67+
68+
protected state: AnyConnectionState = { tag: "disconnected" };
69+
70+
get currentConnectionState(): AnyConnectionState {
71+
return this.state;
72+
}
73+
74+
changeState<Event extends keyof ConnectionManagerEvents, State extends ConnectionManagerEvents[Event][0]>(
75+
event: Event,
76+
newState: State
77+
): State {
78+
this.state = newState;
79+
// TypeScript doesn't seem to be happy with the spread operator and generics
80+
// eslint-disable-next-line
81+
this._events.emit(event, ...([newState] as any));
82+
return newState;
83+
}
84+
85+
setClientName(clientName: string): void {
86+
this.clientName = clientName;
87+
}
88+
89+
abstract connect(connectParams: ConnectParams): Promise<AnyConnectionState>;
90+
91+
abstract disconnect(): Promise<ConnectionStateDisconnected | ConnectionStateErrored>;
92+
}

‎src/common/mcpConnectionManager.ts

Lines changed: 21 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -12,69 +12,18 @@ import type { CompositeLogger } from "./logger.js";
1212
import { LogId } from "./logger.js";
1313
import type { ConnectionInfo } from "@mongosh/arg-parser";
1414
import { generateConnectionInfoFromCliArgs } from "@mongosh/arg-parser";
15-
16-
export interface AtlasClusterConnectionInfo {
17-
username: string;
18-
projectId: string;
19-
clusterName: string;
20-
expiryDate: Date;
21-
}
22-
23-
export interface ConnectionSettings {
24-
connectionString: string;
25-
atlas?: AtlasClusterConnectionInfo;
26-
}
27-
28-
type ConnectionTag = "connected" | "connecting" | "disconnected" | "errored";
29-
type OIDCConnectionAuthType = "oidc-auth-flow" | "oidc-device-flow";
30-
export type ConnectionStringAuthType = "scram" | "ldap" | "kerberos" | OIDCConnectionAuthType | "x.509";
31-
32-
export interface ConnectionState {
33-
tag: ConnectionTag;
34-
connectionStringAuthType?: ConnectionStringAuthType;
35-
connectedAtlasCluster?: AtlasClusterConnectionInfo;
36-
}
37-
38-
export interface ConnectionStateConnected extends ConnectionState {
39-
tag: "connected";
40-
serviceProvider: NodeDriverServiceProvider;
41-
}
42-
43-
export interface ConnectionStateConnecting extends ConnectionState {
44-
tag: "connecting";
45-
serviceProvider: NodeDriverServiceProvider;
46-
oidcConnectionType: OIDCConnectionAuthType;
47-
oidcLoginUrl?: string;
48-
oidcUserCode?: string;
49-
}
50-
51-
export interface ConnectionStateDisconnected extends ConnectionState {
52-
tag: "disconnected";
53-
}
54-
55-
export interface ConnectionStateErrored extends ConnectionState {
56-
tag: "errored";
57-
errorReason: string;
58-
}
59-
60-
export type AnyConnectionState =
61-
| ConnectionStateConnected
62-
| ConnectionStateConnecting
63-
| ConnectionStateDisconnected
64-
| ConnectionStateErrored;
65-
66-
export interface MCPConnectionManagerEvents {
67-
"connection-requested": [AnyConnectionState];
68-
"connection-succeeded": [ConnectionStateConnected];
69-
"connection-timed-out": [ConnectionStateErrored];
70-
"connection-closed": [ConnectionStateDisconnected];
71-
"connection-errored": [ConnectionStateErrored];
72-
}
73-
74-
export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvents> {
75-
private state: AnyConnectionState;
15+
import {
16+
ConnectionManager,
17+
type AnyConnectionState,
18+
type ConnectionStringAuthType,
19+
type OIDCConnectionAuthType,
20+
type ConnectionStateDisconnected,
21+
type ConnectionStateErrored,
22+
type MCPConnectParams,
23+
} from "./connectionManager.js";
24+
25+
export class MCPConnectionManager extends ConnectionManager<MCPConnectParams> {
7626
private deviceId: DeviceId;
77-
private clientName: string;
7827
private bus: EventEmitter;
7928

8029
constructor(
@@ -85,23 +34,15 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
8534
bus?: EventEmitter
8635
) {
8736
super();
88-
8937
this.bus = bus ?? new EventEmitter();
90-
this.state = { tag: "disconnected" };
91-
9238
this.bus.on("mongodb-oidc-plugin:auth-failed", this.onOidcAuthFailed.bind(this));
9339
this.bus.on("mongodb-oidc-plugin:auth-succeeded", this.onOidcAuthSucceeded.bind(this));
94-
9540
this.deviceId = deviceId;
9641
this.clientName = "unknown";
9742
}
9843

99-
setClientName(clientName: string): void {
100-
this.clientName = clientName;
101-
}
102-
103-
async connect(settings: ConnectionSettings): Promise<AnyConnectionState> {
104-
this.emit("connection-requested", this.state);
44+
async connect(connectParams: MCPConnectParams): Promise<AnyConnectionState> {
45+
this._events.emit("connection-requested", this.state);
10546

10647
if (this.state.tag === "connected" || this.state.tag === "connecting") {
10748
await this.disconnect();
@@ -111,22 +52,22 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
11152
let connectionInfo: ConnectionInfo;
11253

11354
try {
114-
settings = { ...settings };
55+
connectParams = { ...connectParams };
11556
const appNameComponents: AppNameComponents = {
11657
appName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
11758
deviceId: this.deviceId.get(),
11859
clientName: this.clientName,
11960
};
12061

121-
settings.connectionString = await setAppNameParamIfMissing({
122-
connectionString: settings.connectionString,
62+
connectParams.connectionString = await setAppNameParamIfMissing({
63+
connectionString: connectParams.connectionString,
12364
components: appNameComponents,
12465
});
12566

12667
connectionInfo = generateConnectionInfoFromCliArgs({
12768
...this.userConfig,
12869
...this.driverOptions,
129-
connectionSpecifier: settings.connectionString,
70+
connectionSpecifier: connectParams.connectionString,
13071
});
13172

13273
if (connectionInfo.driverOptions.oidc) {
@@ -152,7 +93,7 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
15293
this.changeState("connection-errored", {
15394
tag: "errored",
15495
errorReason,
155-
connectedAtlasCluster: settings.atlas,
96+
connectedAtlasCluster: connectParams.atlas,
15697
});
15798
throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, errorReason);
15899
}
@@ -167,7 +108,7 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
167108

168109
return this.changeState("connection-requested", {
169110
tag: "connecting",
170-
connectedAtlasCluster: settings.atlas,
111+
connectedAtlasCluster: connectParams.atlas,
171112
serviceProvider,
172113
connectionStringAuthType: connectionType,
173114
oidcConnectionType: connectionType as OIDCConnectionAuthType,
@@ -178,7 +119,7 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
178119

179120
return this.changeState("connection-succeeded", {
180121
tag: "connected",
181-
connectedAtlasCluster: settings.atlas,
122+
connectedAtlasCluster: connectParams.atlas,
182123
serviceProvider,
183124
connectionStringAuthType: connectionType,
184125
});
@@ -187,7 +128,7 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
187128
this.changeState("connection-errored", {
188129
tag: "errored",
189130
errorReason,
190-
connectedAtlasCluster: settings.atlas,
131+
connectedAtlasCluster: connectParams.atlas,
191132
});
192133
throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, errorReason);
193134
}
@@ -211,21 +152,6 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
211152
return { tag: "disconnected" };
212153
}
213154

214-
get currentConnectionState(): AnyConnectionState {
215-
return this.state;
216-
}
217-
218-
changeState<Event extends keyof MCPConnectionManagerEvents, State extends MCPConnectionManagerEvents[Event][0]>(
219-
event: Event,
220-
newState: State
221-
): State {
222-
this.state = newState;
223-
// TypeScript doesn't seem to be happy with the spread operator and generics
224-
// eslint-disable-next-line
225-
this.emit(event, ...([newState] as any));
226-
return newState;
227-
}
228-
229155
private onOidcAuthFailed(error: unknown): void {
230156
if (this.state.tag === "connecting" && this.state.connectionStringAuthType?.startsWith("oidc")) {
231157
void this.disconnectOnOidcError(error);

‎src/common/session.ts

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,10 @@ import { LogId } from "./logger.js";
77
import EventEmitter from "events";
88
import type {
99
AtlasClusterConnectionInfo,
10-
MCPConnectionManager,
11-
ConnectionSettings,
10+
ConnectionManager,
1211
ConnectionStateConnected,
13-
} from "./mcpConnectionManager.js";
12+
MCPConnectParams,
13+
} from "./connectionManager.js";
1414
import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
1515
import { ErrorCodes, MongoDBError } from "./errors.js";
1616
import type { ExportsManager } from "./exportsManager.js";
@@ -21,7 +21,7 @@ export interface SessionOptions {
2121
apiClientSecret?: string;
2222
logger: CompositeLogger;
2323
exportsManager: ExportsManager;
24-
connectionManager: MCPConnectionManager;
24+
connectionManager: ConnectionManager;
2525
}
2626

2727
export type SessionEvents = {
@@ -34,7 +34,7 @@ export type SessionEvents = {
3434
export class Session extends EventEmitter<SessionEvents> {
3535
readonly sessionId: string = new ObjectId().toString();
3636
readonly exportsManager: ExportsManager;
37-
readonly connectionManager: MCPConnectionManager;
37+
readonly connectionManager: ConnectionManager;
3838
readonly apiClient: ApiClient;
3939
mcpClient?: {
4040
name?: string;
@@ -66,10 +66,14 @@ export class Session extends EventEmitter<SessionEvents> {
6666
this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials }, logger);
6767
this.exportsManager = exportsManager;
6868
this.connectionManager = connectionManager;
69-
this.connectionManager.on("connection-succeeded", () => this.emit("connect"));
70-
this.connectionManager.on("connection-timed-out", (error) => this.emit("connection-error", error.errorReason));
71-
this.connectionManager.on("connection-closed", () => this.emit("disconnect"));
72-
this.connectionManager.on("connection-errored", (error) => this.emit("connection-error", error.errorReason));
69+
this.connectionManager.events.on("connection-succeeded", () => this.emit("connect"));
70+
this.connectionManager.events.on("connection-timed-out", (error) =>
71+
this.emit("connection-error", error.errorReason)
72+
);
73+
this.connectionManager.events.on("connection-closed", () => this.emit("disconnect"));
74+
this.connectionManager.events.on("connection-errored", (error) =>
75+
this.emit("connection-error", error.errorReason)
76+
);
7377
}
7478

7579
setMcpClient(mcpClient: Implementation | undefined): void {
@@ -135,9 +139,9 @@ export class Session extends EventEmitter<SessionEvents> {
135139
this.emit("close");
136140
}
137141

138-
async connectToMongoDB(settings: ConnectionSettings): Promise<void> {
142+
async connectToMongoDB(connectParams: MCPConnectParams): Promise<void> {
139143
try {
140-
await this.connectionManager.connect({ ...settings });
144+
await this.connectionManager.connect({ ...connectParams });
141145
} catch (error: unknown) {
142146
const message = error instanceof Error ? error.message : (error as string);
143147
this.emit("connection-error", message);

‎src/index.ts

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,17 +42,23 @@ import { packageInfo } from "./common/packageInfo.js";
4242
import { StdioRunner } from "./transports/stdio.js";
4343
import { StreamableHttpRunner } from "./transports/streamableHttp.js";
4444
import { systemCA } from "@mongodb-js/devtools-proxy-support";
45+
import type { MCPConnectParams } from "./lib.js";
46+
import type { CreateConnectionManagerFn } from "./transports/base.js";
47+
import { MCPConnectionManager } from "./common/mcpConnectionManager.js";
4548

4649
async function main(): Promise<void> {
4750
systemCA().catch(() => undefined); // load system CA asynchronously as in mongosh
4851

4952
assertHelpMode();
5053
assertVersionMode();
5154

55+
const createConnectionManager: CreateConnectionManagerFn<MCPConnectParams> = ({ logger, deviceId }) =>
56+
new MCPConnectionManager(config, driverOptions, logger, deviceId);
57+
5258
const transportRunner =
5359
config.transport === "stdio"
54-
? new StdioRunner(config, driverOptions)
55-
: new StreamableHttpRunner(config, driverOptions);
60+
? new StdioRunner<MCPConnectParams>(config, createConnectionManager)
61+
: new StreamableHttpRunner<MCPConnectParams>(config, createConnectionManager);
5662
const shutdown = (): void => {
5763
transportRunner.logger.info({
5864
id: LogId.serverCloseRequested,

‎src/lib.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,5 +3,5 @@ export { Telemetry } from "./telemetry/telemetry.js";
33
export { Session, type SessionOptions } from "./common/session.js";
44
export { type UserConfig, defaultUserConfig } from "./common/config.js";
55
export { StreamableHttpRunner } from "./transports/streamableHttp.js";
6-
export { LoggerBase } from "./common/logger.js";
7-
export type { LogPayload, LoggerType, LogLevel } from "./common/logger.js";
6+
export { LoggerBase, CompositeLogger, type LogPayload, type LoggerType, type LogLevel } from "./common/logger.js";
7+
export * from "./common/connectionManager.js";

‎src/tools/atlas/connect/connectCluster.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { generateSecurePassword } from "../../../helpers/generatePassword.js";
66
import { LogId } from "../../../common/logger.js";
77
import { inspectCluster } from "../../../common/atlas/cluster.js";
88
import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
9-
import type { AtlasClusterConnectionInfo } from "../../../common/mcpConnectionManager.js";
9+
import type { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
1010
import { getDefaultRoleFromConfig } from "../../../common/atlas/roles.js";
1111

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

0 commit comments

Comments
 (0)