diff --git a/src/partition/partitioncli.js b/src/partition/partitioncli.js new file mode 100644 index 00000000..d7d2e0ca --- /dev/null +++ b/src/partition/partitioncli.js @@ -0,0 +1,82 @@ +import { PartitionModule } from './partitionmodule'; +import { PartitionMap } from './partitionmap'; +import fs from 'fs'; + +const partitionModule = new PartitionModule(); +partitionModule.cliMode = true; + +function savePartitionMap(filename) { + const reloader = partitionModule.getPartitionMapReloader(); + reloader.once('partitionMapChange', map => { + var toml = 'pool = [\n'; + map.getPool().forEach((poolEntry, i) => { + toml += ` '${poolEntry}'`; + if (i < map.getPool().length - 1) { + toml += ','; + } + + toml += '\n'; + }); + toml += ']\n\n'; + + const partitions = map.getPartitions(); + Object.keys(partitions).forEach(identity => { + partitions[identity].servers.forEach(serverDef => { + toml += `[[partitions.${identity}.servers]]\n`; + toml += `url = '${serverDef.url}'\n`; + toml += `secure = ${serverDef.secure}\n`; + toml += '\n'; + }); + }); + + toml += '[overrides]\n'; + const overrides = map.getOverrides(); + Object.keys(overrides).forEach(channel => { + toml += `${channel} = '${overrides[channel]}'\n`; + }); + + fs.writeFileSync(filename, toml); + console.log(`Wrote partition map to ${filename}`); + process.exit(0); + }); +} + +function loadPartitionMap(filename) { + var newMap; + + try { + newMap = PartitionMap.fromFile(filename); + } catch (error) { + console.error(`Failed to load partition map from ${filename}: ${error}`); + console.error(error.stack); + process.exit(1); + } + + const client = partitionModule.getRedisClientProvider().get(); + client.once('ready', () => { + client.multi() + .set('partitionMap', JSON.stringify(newMap)) + .publish('partitionMap', new Date().toISOString()) + .execAsync() + .then(result => { + console.log(`Result: ${result}`); + console.log(`Published new partition map from ${filename}`); + process.exit(0); + }).catch(error => { + console.error(`Failed to publish partition map: ${error}`); + console.error(error.stack); + process.exit(1); + }); + }); +} + +if (process.argv[2] === 'save') { + savePartitionMap(process.argv[3]); +} else if (process.argv[2] === 'load') { + loadPartitionMap(process.argv[3]); +} else { + console.error('Usage: ' + process.argv[0] + ' ' + process.argv[1] + ' '); + console.error(' "save" downloads the partition map and saves it to the specified file'); + console.error(' "load" loads the partition map from the specified file and publishes it'); + process.exit(1); +} diff --git a/src/partition/partitionmodule.js b/src/partition/partitionmodule.js index 42639b53..d698b012 100644 --- a/src/partition/partitionmodule.js +++ b/src/partition/partitionmodule.js @@ -15,6 +15,7 @@ const PARTITION_CONFIG_PATH = path.resolve(__dirname, '..', '..', 'conf', class PartitionModule { constructor() { this.initConfig(); + this.cliMode = false; } onReady() { @@ -63,7 +64,9 @@ class PartitionModule { reloader.getPartitionMap()); reloader.on('partitionMapChange', newMap => { this.partitionDecider.setPartitionMap(newMap); - require('../server').getServer().handlePartitionMapChange(); + if (!this.cliMode) { + require('../server').getServer().handlePartitionMapChange(); + } }); }