Skip to content

Commit bf532da

Browse files
claabsclaabs
committed
Make use of extendable connections
1 parent d840b61 commit bf532da

File tree

3 files changed

+83
-51
lines changed

3 files changed

+83
-51
lines changed

‎README.md

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ Give it an array of strings, and it will output a randomly generated string.
1616

1717
This rewrite was created for the Discord bot [markov-discord](https://github.com/claabs/markov-discord).
1818

19-
2019
## Prerequisites
2120

2221
- NodeJS 10+
@@ -36,11 +35,17 @@ import Markov from 'markov-strings-db';
3635

3736
const data = [/* insert a few hundreds/thousands sentences here */];
3837

38+
3939
// Instantiate the Markov generator
4040
const markov = new Markov({ options: { stateSize: 2 }});
4141

42-
// Connect to the database. This is required for anything to work.
43-
await markov.connect();
42+
// If you have your own database you'd like to combine with Markov's, make sure to extend your connection
43+
const connectionOptions = Markov.extendConnectionOptions();
44+
// Required: create a connection before using a markov instance. You only need to do this once.
45+
const connection = await createConnection(connectionOptions);
46+
47+
// If you have a non-default connection (you probably don't), pass it in here. Otherwise, setup() is called implicitly on any async function.
48+
await markov.setup(connection);
4449

4550
// Add data for the generator
4651
await markov.addData(data)

‎src/index.ts

Lines changed: 48 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/* eslint-disable no-await-in-loop */
2-
import { Connection, ConnectionOptions, createConnection, getConnectionOptions, In } from 'typeorm';
2+
import { Connection, ConnectionOptions, EntitySchema, getConnectionOptions, In } from 'typeorm';
33
import { MarkovCorpusEntry } from './entity/MarkovCorpusEntry';
44
import { MarkovFragment } from './entity/MarkovFragment';
55
import { MarkovInputData } from './entity/MarkovInputData';
@@ -8,6 +8,14 @@ import { MarkovRoot } from './entity/MarkovRoot';
88
import { Importer } from './importer';
99
import { MarkovImportExport as MarkovV3ImportExport } from './v3-types';
1010

11+
const ALL_ENTITIES = [
12+
MarkovCorpusEntry,
13+
MarkovRoot,
14+
MarkovOptions,
15+
MarkovInputData,
16+
MarkovFragment,
17+
];
18+
1119
/**
1220
* Data to build the Markov instance
1321
*/
@@ -100,14 +108,36 @@ export default class Markov {
100108

101109
public options: MarkovOptions | MarkovDataMembers;
102110

103-
public connection: Connection;
104-
105111
public id: string;
106112

107113
private defaultOptions: MarkovDataMembers = {
108114
stateSize: 2,
109115
};
110116

117+
/**
118+
* If you're connecting the Markov DB with your parent project's DB, you'll need to extend it to add the Markov entities to the connection.
119+
* @param connectionOptions Your TypeORM connection options. If not provided, it will load the ormconfig from a file.
120+
* @returns ConnectionOptions that should be passed into a createConnection() call.
121+
*/
122+
public static async extendConnectionOptions(
123+
connectionOptions?: ConnectionOptions
124+
): Promise<ConnectionOptions> {
125+
let baseConnectionOpts: ConnectionOptions;
126+
if (connectionOptions) {
127+
baseConnectionOpts = connectionOptions;
128+
} else {
129+
baseConnectionOpts = await getConnectionOptions();
130+
}
131+
132+
// eslint-disable-next-line @typescript-eslint/ban-types, @typescript-eslint/no-explicit-any
133+
const appendedEntities: (Function | string | EntitySchema<any>)[] = ALL_ENTITIES;
134+
if (baseConnectionOpts.entities) appendedEntities.push(...baseConnectionOpts.entities);
135+
return {
136+
...baseConnectionOpts,
137+
entities: appendedEntities,
138+
};
139+
}
140+
111141
/**
112142
* Creates an instance of Markov generator.
113143
*/
@@ -119,19 +149,12 @@ export default class Markov {
119149
}
120150

121151
/**
122-
* Connects this instance to your database. You should setup an ormconfig, or pass it in as an object.
152+
* If you have a non-default connection (you probably don't), pass it in here.
153+
* Otherwise, setup() is called implicitly on any async function.
154+
* @param connection A non-default connection to be used by Markov's entities
123155
*/
124-
public async connect(connectionOptions?: ConnectionOptions): Promise<Connection> {
125-
let baseConnectionOpts: ConnectionOptions;
126-
if (connectionOptions) {
127-
baseConnectionOpts = connectionOptions;
128-
} else {
129-
baseConnectionOpts = await getConnectionOptions();
130-
}
131-
this.connection = await createConnection({
132-
...baseConnectionOpts,
133-
entities: [MarkovCorpusEntry, MarkovRoot, MarkovOptions, MarkovInputData, MarkovFragment],
134-
});
156+
public async setup(connection?: Connection): Promise<void> {
157+
if (connection) ALL_ENTITIES.forEach((e) => e.useConnection(connection));
135158

136159
let db = await MarkovRoot.findOne({
137160
id: this.id,
@@ -151,13 +174,10 @@ export default class Markov {
151174
this.db = db;
152175
this.id = this.db.id;
153176
this.options = options;
154-
return this.connection;
155177
}
156178

157-
public async disconnect(): Promise<void> {
158-
if (this.connection) {
159-
await this.connection.close();
160-
}
179+
private async ensureSetup(): Promise<void> {
180+
if (!this.db) await this.setup();
161181
}
162182

163183
/**
@@ -183,12 +203,13 @@ export default class Markov {
183203
* Supports imports from markov-strings v3, as well as exports from this version.
184204
*/
185205
public async import(data: MarkovRoot | MarkovV3ImportExport): Promise<void> {
206+
await this.ensureSetup();
186207
if ('id' in data) {
187-
const options = MarkovOptions.create(data.options);
188-
await MarkovOptions.save(options);
189208
this.db = new MarkovRoot();
190209
this.db.id = this.id;
210+
const options = MarkovOptions.create(data.options);
191211
this.db.options = options;
212+
await MarkovOptions.save(options);
192213

193214
const importer = new Importer(this.db);
194215
const startWords = await importer.saveImportFragments(data.startWords, 'startWordMarkov');
@@ -228,6 +249,7 @@ export default class Markov {
228249
* Exports all the data in the database associated with this Markov instance as a JSON object.
229250
*/
230251
public async export(): Promise<MarkovRoot> {
252+
await this.ensureSetup();
231253
const db = await MarkovRoot.findOneOrFail({
232254
where: { id: this.id },
233255
relations: [
@@ -251,6 +273,7 @@ export default class Markov {
251273
* It's possible to store custom JSON-like data of your choice next to the string by passing in objects of `{ string: 'foo', custom: 'attachment' }`. This data will be returned in the `refs` of the result.
252274
*/
253275
public async addData(rawData: AddDataProps[] | string[]) {
276+
await this.ensureSetup();
254277
// Format data if necessary
255278
let input: AddDataProps[] = [];
256279
if (typeof rawData[0] === 'string') {
@@ -270,6 +293,7 @@ export default class Markov {
270293
* Builds the corpus. You must call this before generating sentences.
271294
*/
272295
private async buildCorpus(data: AddDataProps[]): Promise<void> {
296+
await this.ensureSetup();
273297
const { options } = this.db;
274298

275299
// Loop through all sentences
@@ -412,6 +436,7 @@ export default class Markov {
412436
* @param rawData A list of full strings
413437
*/
414438
public async removeData(rawData: string[]): Promise<void> {
439+
await this.ensureSetup();
415440
const inputData = await MarkovInputData.find({
416441
relations: ['fragment', 'fragment.corpusEntry'],
417442
where: [
@@ -447,15 +472,14 @@ export default class Markov {
447472
public async generate<CustomData = any>( // eslint-disable-line @typescript-eslint/no-explicit-any
448473
options?: MarkovGenerateOptions<CustomData>
449474
): Promise<MarkovResult<CustomData>> {
450-
// const corpusSize = await CorpusEntry.count({markov: this.db});
475+
await this.ensureSetup();
451476
const corpusSize = await MarkovCorpusEntry.count({ markov: this.db });
452477
if (corpusSize <= 0) {
453478
throw new Error(
454479
'Corpus is empty. There is either no data, or the data is not sufficient to create markov chains.'
455480
);
456481
}
457482

458-
// const corpus = cloneDeep(this.corpus)
459483
const maxTries = options?.maxTries ? options.maxTries : 10;
460484

461485
let tries: number;

‎test/test.ts

Lines changed: 27 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import path from 'path';
22
import { readFileSync } from 'fs';
3+
import { Connection, createConnection } from 'typeorm';
34
import Markov, { AddDataProps, MarkovResult } from '../src/index';
45
import { MarkovCorpusEntry } from '../src/entity/MarkovCorpusEntry';
56
import { MarkovFragment } from '../src/entity/MarkovFragment';
@@ -18,6 +19,8 @@ const data = [
1819

1920
jest.setTimeout(1000000);
2021
describe('Markov class', () => {
22+
let connection: Connection;
23+
2124
describe('Constructor', () => {
2225
it('should have a default stateSize', () => {
2326
const markov = new Markov();
@@ -30,27 +33,27 @@ describe('Markov class', () => {
3033
});
3134

3235
it('should persist options in the database', async () => {
36+
connection = await createConnection();
3337
const markov = new Markov({ options: { stateSize: 3 } });
34-
await markov.connect();
35-
await markov.disconnect();
38+
await markov.setup();
3639
const markov2 = new Markov();
37-
await markov2.connect();
40+
await markov2.setup();
3841
expect(markov2.options.stateSize).toBe(3);
39-
await markov2.connection.dropDatabase();
40-
await markov2.disconnect();
42+
await connection.dropDatabase();
43+
await connection.close();
4144
});
4245
});
4346

4447
describe('Adding data', () => {
4548
let markov: Markov;
4649
beforeEach(async () => {
4750
markov = new Markov();
48-
await markov.connect();
51+
connection = await createConnection();
4952
});
5053

5154
afterEach(async () => {
52-
await markov.connection.dropDatabase();
53-
await markov.disconnect();
55+
await connection.dropDatabase();
56+
await connection.close();
5457
});
5558

5659
it('should build corpus', async () => {
@@ -82,13 +85,13 @@ describe('Markov class', () => {
8285
let markov: Markov;
8386
beforeEach(async () => {
8487
markov = new Markov();
85-
await markov.connect();
88+
connection = await createConnection();
8689
await markov.addData(data);
8790
});
8891

8992
afterEach(async () => {
90-
await markov.connection.dropDatabase();
91-
await markov.disconnect();
93+
await connection.dropDatabase();
94+
await connection.close();
9295
});
9396

9497
describe('The startWords array', () => {
@@ -175,13 +178,13 @@ describe('Markov class', () => {
175178
describe('Import/Export', () => {
176179
let markov: Markov;
177180
afterEach(async () => {
178-
await markov.connection.dropDatabase();
179-
await markov.disconnect();
181+
await connection.dropDatabase();
182+
await connection.close();
180183
});
181184

182185
it('should export the original database values', async () => {
183186
markov = new Markov();
184-
await markov.connect();
187+
connection = await createConnection();
185188
await markov.addData(data);
186189

187190
const exported = await markov.export();
@@ -194,7 +197,7 @@ describe('Markov class', () => {
194197
describe('Import v3 data', () => {
195198
it('onto fresh DB', async () => {
196199
markov = new Markov();
197-
await markov.connect();
200+
connection = await createConnection();
198201
let count = await MarkovCorpusEntry.count({
199202
markov: markov.db,
200203
});
@@ -215,7 +218,7 @@ describe('Markov class', () => {
215218

216219
it('should overwrite original values', async () => {
217220
markov = new Markov();
218-
await markov.connect();
221+
connection = await createConnection();
219222
await markov.addData(data);
220223

221224
const v3Import = JSON.parse(readFileSync(path.join(__dirname, 'v3-export.json'), 'utf8'));
@@ -231,7 +234,7 @@ describe('Markov class', () => {
231234
describe('Import v4 data', () => {
232235
it('onto fresh DB', async () => {
233236
markov = new Markov();
234-
await markov.connect();
237+
connection = await createConnection();
235238
let count = await MarkovCorpusEntry.count({
236239
markov: markov.db,
237240
});
@@ -252,7 +255,7 @@ describe('Markov class', () => {
252255

253256
it('should overwrite original values', async () => {
254257
markov = new Markov();
255-
await markov.connect();
258+
connection = await createConnection();
256259
await markov.addData(data);
257260

258261
const v4Import = JSON.parse(readFileSync(path.join(__dirname, 'v4-export.json'), 'utf8'));
@@ -271,12 +274,12 @@ describe('Markov class', () => {
271274
describe('With no data', () => {
272275
beforeEach(async () => {
273276
markov = new Markov();
274-
await markov.connect();
277+
connection = await createConnection();
275278
});
276279

277280
afterEach(async () => {
278-
await markov.connection.dropDatabase();
279-
await markov.disconnect();
281+
await connection.dropDatabase();
282+
await connection.close();
280283
});
281284

282285
it('should throw an error if the corpus is not built', async () => {
@@ -314,13 +317,13 @@ describe('Markov class', () => {
314317
describe('With data', () => {
315318
beforeEach(async () => {
316319
markov = new Markov();
317-
await markov.connect();
320+
connection = await createConnection();
318321
await markov.addData(data);
319322
});
320323

321324
afterEach(async () => {
322-
await markov.connection.dropDatabase();
323-
await markov.disconnect();
325+
await connection.dropDatabase();
326+
await connection.close();
324327
});
325328

326329
it('should return a result if under the tries limit', async () => {

0 commit comments

Comments
 (0)