396 lines
11 KiB
TypeScript
396 lines
11 KiB
TypeScript
import {
|
|
Collection,
|
|
Db,
|
|
Document,
|
|
MongoClient,
|
|
MongoClientOptions,
|
|
OptionalUnlessRequiredId,
|
|
WithId,
|
|
} from 'mongodb';
|
|
import * as yup from 'yup';
|
|
import { mongodbConfig } from '@stock-bot/config';
|
|
import { getLogger } from '@stock-bot/logger';
|
|
import { MongoDBHealthMonitor } from './health';
|
|
import { schemaMap } from './schemas';
|
|
import type {
|
|
AnalystReport,
|
|
CollectionNames,
|
|
DocumentBase,
|
|
EarningsTranscript,
|
|
MongoDBClientConfig,
|
|
MongoDBConnectionOptions,
|
|
NewsArticle,
|
|
RawDocument,
|
|
SecFiling,
|
|
SentimentData,
|
|
} from './types';
|
|
|
|
/**
|
|
* MongoDB Client for Stock Bot
|
|
*
|
|
* Provides type-safe access to MongoDB collections with built-in
|
|
* health monitoring, connection pooling, and schema validation.
|
|
*/
|
|
export class MongoDBClient {
|
|
private client: MongoClient | null = null;
|
|
private db: Db | null = null;
|
|
private readonly config: MongoDBClientConfig;
|
|
private readonly options: MongoDBConnectionOptions;
|
|
private readonly logger: ReturnType<typeof getLogger>;
|
|
private readonly healthMonitor: MongoDBHealthMonitor;
|
|
private isConnected = false;
|
|
|
|
constructor(config?: Partial<MongoDBClientConfig>, options?: MongoDBConnectionOptions) {
|
|
this.config = this.buildConfig(config);
|
|
this.options = {
|
|
retryAttempts: 3,
|
|
retryDelay: 1000,
|
|
healthCheckInterval: 30000,
|
|
...options,
|
|
};
|
|
|
|
this.logger = getLogger('mongodb-client');
|
|
this.healthMonitor = new MongoDBHealthMonitor(this);
|
|
}
|
|
|
|
/**
|
|
* Connect to MongoDB
|
|
*/
|
|
async connect(): Promise<void> {
|
|
if (this.isConnected && this.client) {
|
|
return;
|
|
}
|
|
|
|
const uri = this.buildConnectionUri();
|
|
const clientOptions = this.buildClientOptions();
|
|
|
|
let lastError: Error | null = null;
|
|
|
|
for (let attempt = 1; attempt <= this.options.retryAttempts!; attempt++) {
|
|
try {
|
|
this.logger.info(
|
|
`Connecting to MongoDB (attempt ${attempt}/${this.options.retryAttempts})...`
|
|
);
|
|
|
|
this.client = new MongoClient(uri, clientOptions);
|
|
await this.client.connect();
|
|
|
|
// Test the connection
|
|
await this.client.db(this.config.database).admin().ping();
|
|
|
|
this.db = this.client.db(this.config.database);
|
|
this.isConnected = true;
|
|
|
|
this.logger.info('Successfully connected to MongoDB');
|
|
|
|
// Start health monitoring
|
|
this.healthMonitor.start();
|
|
|
|
return;
|
|
} catch (error) {
|
|
lastError = error as Error;
|
|
this.logger.error(`MongoDB connection attempt ${attempt} failed:`, error);
|
|
|
|
if (this.client) {
|
|
await this.client.close();
|
|
this.client = null;
|
|
}
|
|
|
|
if (attempt < this.options.retryAttempts!) {
|
|
await this.delay(this.options.retryDelay! * attempt);
|
|
}
|
|
}
|
|
}
|
|
|
|
throw new Error(
|
|
`Failed to connect to MongoDB after ${this.options.retryAttempts} attempts: ${lastError?.message}`
|
|
);
|
|
}
|
|
|
|
/**
|
|
* Disconnect from MongoDB
|
|
*/
|
|
async disconnect(): Promise<void> {
|
|
if (!this.client) {
|
|
return;
|
|
}
|
|
|
|
try {
|
|
this.healthMonitor.stop();
|
|
await this.client.close();
|
|
this.isConnected = false;
|
|
this.client = null;
|
|
this.db = null;
|
|
this.logger.info('Disconnected from MongoDB');
|
|
} catch (error) {
|
|
this.logger.error('Error disconnecting from MongoDB:', error);
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Get a typed collection
|
|
*/
|
|
getCollection<T extends DocumentBase>(name: CollectionNames): Collection<T> {
|
|
if (!this.db) {
|
|
throw new Error('MongoDB client not connected');
|
|
}
|
|
return this.db.collection<T>(name);
|
|
}
|
|
|
|
/**
|
|
* Insert a document with validation
|
|
*/
|
|
async insertOne<T extends DocumentBase>(
|
|
collectionName: CollectionNames,
|
|
document: Omit<T, '_id' | 'created_at' | 'updated_at'> &
|
|
Partial<Pick<T, 'created_at' | 'updated_at'>>
|
|
): Promise<T> {
|
|
const collection = this.getCollection<T>(collectionName);
|
|
|
|
// Add timestamps
|
|
const now = new Date();
|
|
const docWithTimestamps = {
|
|
...document,
|
|
created_at: document.created_at || now,
|
|
updated_at: now,
|
|
} as T; // Validate document if schema exists
|
|
if (collectionName in schemaMap) {
|
|
try {
|
|
(schemaMap as any)[collectionName].validateSync(docWithTimestamps);
|
|
} catch (error) {
|
|
if (error instanceof yup.ValidationError) {
|
|
this.logger.error(`Document validation failed for ${collectionName}:`, error.errors);
|
|
throw new Error(`Document validation failed: ${error.errors?.map(e => e).join(', ')}`);
|
|
}
|
|
throw error;
|
|
}
|
|
}
|
|
const result = await collection.insertOne(docWithTimestamps as OptionalUnlessRequiredId<T>);
|
|
return { ...docWithTimestamps, _id: result.insertedId } as T;
|
|
}
|
|
|
|
/**
|
|
* Update a document with validation
|
|
*/
|
|
async updateOne<T extends DocumentBase>(
|
|
collectionName: CollectionNames,
|
|
filter: any,
|
|
update: Partial<T>
|
|
): Promise<boolean> {
|
|
const collection = this.getCollection<T>(collectionName);
|
|
|
|
// Add updated timestamp
|
|
const updateWithTimestamp = {
|
|
...update,
|
|
updated_at: new Date(),
|
|
};
|
|
|
|
const result = await collection.updateOne(filter, { $set: updateWithTimestamp });
|
|
return result.modifiedCount > 0;
|
|
}
|
|
/**
|
|
* Find documents with optional validation
|
|
*/
|
|
async find<T extends DocumentBase>(
|
|
collectionName: CollectionNames,
|
|
filter: any = {},
|
|
options: any = {}
|
|
): Promise<T[]> {
|
|
const collection = this.getCollection<T>(collectionName);
|
|
return (await collection.find(filter, options).toArray()) as T[];
|
|
}
|
|
|
|
/**
|
|
* Find one document
|
|
*/
|
|
async findOne<T extends DocumentBase>(
|
|
collectionName: CollectionNames,
|
|
filter: any
|
|
): Promise<T | null> {
|
|
const collection = this.getCollection<T>(collectionName);
|
|
return (await collection.findOne(filter)) as T | null;
|
|
}
|
|
|
|
/**
|
|
* Aggregate with type safety
|
|
*/
|
|
async aggregate<T extends DocumentBase>(
|
|
collectionName: CollectionNames,
|
|
pipeline: any[]
|
|
): Promise<T[]> {
|
|
const collection = this.getCollection<T>(collectionName);
|
|
return await collection.aggregate<T>(pipeline).toArray();
|
|
}
|
|
|
|
/**
|
|
* Count documents
|
|
*/
|
|
async countDocuments(collectionName: CollectionNames, filter: any = {}): Promise<number> {
|
|
const collection = this.getCollection(collectionName);
|
|
return await collection.countDocuments(filter);
|
|
}
|
|
|
|
/**
|
|
* Create indexes for better performance
|
|
*/
|
|
async createIndexes(): Promise<void> {
|
|
if (!this.db) {
|
|
throw new Error('MongoDB client not connected');
|
|
}
|
|
|
|
try {
|
|
// Sentiment data indexes
|
|
await this.db
|
|
.collection('sentiment_data')
|
|
.createIndexes([
|
|
{ key: { symbol: 1, timestamp: -1 } },
|
|
{ key: { sentiment_label: 1 } },
|
|
{ key: { source_type: 1 } },
|
|
{ key: { created_at: -1 } },
|
|
]);
|
|
|
|
// News articles indexes
|
|
await this.db
|
|
.collection('news_articles')
|
|
.createIndexes([
|
|
{ key: { symbols: 1, published_date: -1 } },
|
|
{ key: { publication: 1 } },
|
|
{ key: { categories: 1 } },
|
|
{ key: { created_at: -1 } },
|
|
]);
|
|
|
|
// SEC filings indexes
|
|
await this.db
|
|
.collection('sec_filings')
|
|
.createIndexes([
|
|
{ key: { symbols: 1, filing_date: -1 } },
|
|
{ key: { filing_type: 1 } },
|
|
{ key: { cik: 1 } },
|
|
{ key: { created_at: -1 } },
|
|
]); // Raw documents indexes
|
|
await this.db.collection('raw_documents').createIndex({ content_hash: 1 }, { unique: true });
|
|
await this.db
|
|
.collection('raw_documents')
|
|
.createIndexes([
|
|
{ key: { processing_status: 1 } },
|
|
{ key: { document_type: 1 } },
|
|
{ key: { created_at: -1 } },
|
|
]);
|
|
|
|
this.logger.info('MongoDB indexes created successfully');
|
|
} catch (error) {
|
|
this.logger.error('Error creating MongoDB indexes:', error);
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Get database statistics
|
|
*/
|
|
async getStats(): Promise<any> {
|
|
if (!this.db) {
|
|
throw new Error('MongoDB client not connected');
|
|
}
|
|
return await this.db.stats();
|
|
}
|
|
|
|
/**
|
|
* Check if client is connected
|
|
*/
|
|
get connected(): boolean {
|
|
return this.isConnected && !!this.client;
|
|
}
|
|
|
|
/**
|
|
* Get the underlying MongoDB client
|
|
*/
|
|
get mongoClient(): MongoClient | null {
|
|
return this.client;
|
|
}
|
|
|
|
/**
|
|
* Get the database instance
|
|
*/
|
|
get database(): Db | null {
|
|
return this.db;
|
|
}
|
|
|
|
private buildConfig(config?: Partial<MongoDBClientConfig>): MongoDBClientConfig {
|
|
return {
|
|
host: config?.host || mongodbConfig.MONGODB_HOST,
|
|
port: config?.port || mongodbConfig.MONGODB_PORT,
|
|
database: config?.database || mongodbConfig.MONGODB_DATABASE,
|
|
username: config?.username || mongodbConfig.MONGODB_USERNAME,
|
|
password: config?.password || mongodbConfig.MONGODB_PASSWORD,
|
|
authSource: config?.authSource || mongodbConfig.MONGODB_AUTH_SOURCE,
|
|
uri: config?.uri || mongodbConfig.MONGODB_URI,
|
|
poolSettings: {
|
|
maxPoolSize: mongodbConfig.MONGODB_MAX_POOL_SIZE,
|
|
minPoolSize: mongodbConfig.MONGODB_MIN_POOL_SIZE,
|
|
maxIdleTime: mongodbConfig.MONGODB_MAX_IDLE_TIME,
|
|
...config?.poolSettings,
|
|
},
|
|
timeouts: {
|
|
connectTimeout: mongodbConfig.MONGODB_CONNECT_TIMEOUT,
|
|
socketTimeout: mongodbConfig.MONGODB_SOCKET_TIMEOUT,
|
|
serverSelectionTimeout: mongodbConfig.MONGODB_SERVER_SELECTION_TIMEOUT,
|
|
...config?.timeouts,
|
|
},
|
|
tls: {
|
|
enabled: mongodbConfig.MONGODB_TLS,
|
|
insecure: mongodbConfig.MONGODB_TLS_INSECURE,
|
|
caFile: mongodbConfig.MONGODB_TLS_CA_FILE,
|
|
...config?.tls,
|
|
},
|
|
options: {
|
|
retryWrites: mongodbConfig.MONGODB_RETRY_WRITES,
|
|
journal: mongodbConfig.MONGODB_JOURNAL,
|
|
readPreference: mongodbConfig.MONGODB_READ_PREFERENCE as any,
|
|
writeConcern: mongodbConfig.MONGODB_WRITE_CONCERN,
|
|
...config?.options,
|
|
},
|
|
};
|
|
}
|
|
|
|
private buildConnectionUri(): string {
|
|
if (this.config.uri) {
|
|
return this.config.uri;
|
|
}
|
|
|
|
const { host, port, username, password, database, authSource } = this.config;
|
|
const auth = username && password ? `${username}:${password}@` : '';
|
|
const authDb = authSource ? `?authSource=${authSource}` : '';
|
|
|
|
return `mongodb://${auth}${host}:${port}/${database}${authDb}`;
|
|
}
|
|
|
|
private buildClientOptions(): MongoClientOptions {
|
|
return {
|
|
maxPoolSize: this.config.poolSettings?.maxPoolSize,
|
|
minPoolSize: this.config.poolSettings?.minPoolSize,
|
|
maxIdleTimeMS: this.config.poolSettings?.maxIdleTime,
|
|
connectTimeoutMS: this.config.timeouts?.connectTimeout,
|
|
socketTimeoutMS: this.config.timeouts?.socketTimeout,
|
|
serverSelectionTimeoutMS: this.config.timeouts?.serverSelectionTimeout,
|
|
retryWrites: this.config.options?.retryWrites,
|
|
journal: this.config.options?.journal,
|
|
readPreference: this.config.options?.readPreference,
|
|
writeConcern: this.config.options?.writeConcern
|
|
? {
|
|
w:
|
|
this.config.options.writeConcern === 'majority'
|
|
? ('majority' as const)
|
|
: parseInt(this.config.options.writeConcern, 10) || 1,
|
|
}
|
|
: undefined,
|
|
tls: this.config.tls?.enabled,
|
|
tlsInsecure: this.config.tls?.insecure,
|
|
tlsCAFile: this.config.tls?.caFile,
|
|
};
|
|
}
|
|
|
|
private delay(ms: number): Promise<void> {
|
|
return new Promise(resolve => setTimeout(resolve, ms));
|
|
}
|
|
}
|