Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1#!/usr/bin/env python 

2# encoding: utf-8 

3""" 

4*Toolset to setup the main function for a cl-util* 

5 

6:Author: 

7 David Young 

8 

9:Date Created: 

10 April 16, 2014 

11""" 

12from __future__ import print_function 

13from __future__ import absolute_import 

14################# GLOBAL IMPORTS #################### 

15from builtins import object 

16import sys 

17import os 

18import yaml 

19try: 

20 yaml.warnings({'YAMLLoadWarning': False}) 

21except: 

22 pass 

23from collections import OrderedDict 

24import shutil 

25from subprocess import Popen, PIPE, STDOUT 

26from . import logs as dl 

27import time 

28from docopt import docopt 

29import psutil 

30try: 

31 from StringIO import StringIO 

32except ImportError: 

33 from io import StringIO 

34from os.path import expanduser 

35 

36################################################################### 

37# CLASSES # 

38################################################################### 

39 

40 

41class tools(object): 

42 """ 

43 *common setup methods & attributes of the main function in cl-util* 

44 

45 **Key Arguments:** 

46 - ``dbConn`` -- mysql database connection 

47 - ``arguments`` -- the arguments read in from the command-line 

48 - ``docString`` -- pass the docstring from the host module so that docopt can work on the usage text to generate the required arguments 

49 - ``logLevel`` -- the level of the logger required. Default *DEBUG*. [DEBUG|INFO|WARNING|ERROR|CRITICAL] 

50 - ``options_first`` -- options come before commands in CL usage. Default *False*. 

51 - ``projectName`` -- the name of the project, used to create a default settings file in ``~/.config/projectName/projectName.yaml``. Default *False*. 

52 - ``distributionName`` -- the distribution name if different from the projectName (i.e. if the package is called by anohter name on PyPi). Default *False* 

53 - ``tunnel`` -- will setup a ssh tunnel (if the settings are found in the settings file). Default *False*. 

54 - ``defaultSettingsFile`` -- if no settings file is passed via the doc-string use the default settings file in ``~/.config/projectName/projectName.yaml`` (don't have to clutter command-line with settings) 

55 

56 **Usage:** 

57 

58 Add this to the ``__main__`` function of your command-line module 

59 

60 .. code-block:: python  

61 

62 # setup the command-line util settings 

63 from fundamentals import tools 

64 su = tools( 

65 arguments=arguments, 

66 docString=__doc__, 

67 logLevel="DEBUG", 

68 options_first=False, 

69 projectName="myprojectName" 

70 ) 

71 arguments, settings, log, dbConn = su.setup() 

72 

73 Here is a template settings file content you could use: 

74 

75 .. code-block:: yaml 

76 

77 version: 1 

78 database settings: 

79 db: unit_tests 

80 host: localhost 

81 user: utuser 

82 password: utpass 

83 tunnel: true 

84 

85 # SSH TUNNEL - if a tunnel is required to connect to the database(s) then add setup here 

86 # Note only one tunnel is setup - may need to change this to 2 tunnels in the future if  

87 # code, static catalogue database and transient database are all on seperate machines. 

88 ssh tunnel: 

89 remote user: username 

90 remote ip: mydomain.co.uk 

91 remote datbase host: mydatabaseName 

92 port: 9002 

93 

94 logging settings: 

95 formatters: 

96 file_style: 

97 format: '* %(asctime)s - %(name)s - %(levelname)s (%(pathname)s > %(funcName)s > %(lineno)d) - %(message)s ' 

98 datefmt: '%Y/%m/%d %H:%M:%S' 

99 console_style: 

100 format: '* %(asctime)s - %(levelname)s: %(pathname)s:%(funcName)s:%(lineno)d > %(message)s' 

101 datefmt: '%H:%M:%S' 

102 html_style: 

103 format: '<div id="row" class="%(levelname)s"><span class="date">%(asctime)s</span> <span class="label">file:</span><span class="filename">%(filename)s</span> <span class="label">method:</span><span class="funcName">%(funcName)s</span> <span class="label">line#:</span><span class="lineno">%(lineno)d</span> <span class="pathname">%(pathname)s</span> <div class="right"><span class="message">%(message)s</span><span class="levelname">%(levelname)s</span></div></div>' 

104 datefmt: '%Y-%m-%d <span class= "time">%H:%M <span class= "seconds">%Ss</span></span>' 

105 handlers: 

106 console: 

107 class: logging.StreamHandler 

108 level: DEBUG 

109 formatter: console_style 

110 stream: ext://sys.stdout 

111 file: 

112 class: logging.handlers.GroupWriteRotatingFileHandler 

113 level: WARNING 

114 formatter: file_style 

115 filename: /Users/Dave/.config/myprojectName/myprojectName.log 

116 mode: w+ 

117 maxBytes: 102400 

118 backupCount: 1 

119 root: 

120 level: WARNING 

121 handlers: [file,console] 

122 """ 

123 # Initialisation 

124 

125 def __init__( 

126 self, 

127 arguments, 

128 docString, 

129 logLevel="WARNING", 

130 options_first=False, 

131 projectName=False, 

132 distributionName=False, 

133 orderedSettings=False, 

134 defaultSettingsFile=False 

135 ): 

136 self.arguments = arguments 

137 self.docString = docString 

138 self.logLevel = logLevel 

139 

140 if not distributionName: 

141 distributionName = projectName 

142 

143 version = '0.0.1' 

144 try: 

145 import pkg_resources 

146 version = pkg_resources.get_distribution(distributionName).version 

147 except: 

148 pass 

149 

150 ## ACTIONS BASED ON WHICH ARGUMENTS ARE RECIEVED ## 

151 # PRINT COMMAND-LINE USAGE IF NO ARGUMENTS PASSED 

152 if arguments == None: 

153 arguments = docopt(docString, version="v" + version, 

154 options_first=options_first) 

155 self.arguments = arguments 

156 

157 # BUILD A STRING FOR THE PROCESS TO MATCH RUNNING PROCESSES AGAINST 

158 lockname = "".join(sys.argv) 

159 

160 # TEST IF THE PROCESS IS ALREADY RUNNING WITH THE SAME ARGUMENTS (e.g. 

161 # FROM CRON) - QUIT IF MATCH FOUND 

162 for q in psutil.process_iter(): 

163 try: 

164 this = q.cmdline() 

165 except: 

166 continue 

167 

168 test = "".join(this[1:]) 

169 if q.pid != os.getpid() and lockname == test and "--reload" not in test: 

170 thisId = q.pid 

171 print("This command is already running (see PID %(thisId)s)" % locals()) 

172 sys.exit(0) 

173 

174 try: 

175 if "tests.test" in arguments["<pathToSettingsFile>"]: 

176 del arguments["<pathToSettingsFile>"] 

177 except: 

178 pass 

179 

180 if defaultSettingsFile and "settingsFile" not in arguments and os.path.exists(os.getenv( 

181 "HOME") + "/.config/%(projectName)s/%(projectName)s.yaml" % locals()): 

182 arguments["settingsFile"] = settingsFile = os.getenv( 

183 "HOME") + "/.config/%(projectName)s/%(projectName)s.yaml" % locals() 

184 

185 # UNPACK SETTINGS 

186 stream = False 

187 if "<settingsFile>" in arguments and arguments["<settingsFile>"]: 

188 stream = open(arguments["<settingsFile>"], 'r') 

189 elif "<pathToSettingsFile>" in arguments and arguments["<pathToSettingsFile>"]: 

190 stream = open(arguments["<pathToSettingsFile>"], 'r') 

191 elif "--settingsFile" in arguments and arguments["--settingsFile"]: 

192 stream = open(arguments["--settingsFile"], 'r') 

193 elif "--settings" in arguments and arguments["--settings"]: 

194 stream = open(arguments["--settings"], 'r') 

195 elif "pathToSettingsFile" in arguments and arguments["pathToSettingsFile"]: 

196 stream = open(arguments["pathToSettingsFile"], 'r') 

197 elif "settingsFile" in arguments and arguments["settingsFile"]: 

198 stream = open(arguments["settingsFile"], 'r') 

199 elif ("settingsFile" in arguments and arguments["settingsFile"] == None) or ("<pathToSettingsFile>" in arguments and arguments["<pathToSettingsFile>"] == None) or ("--settings" in arguments and arguments["--settings"] == None) or ("pathToSettingsFile" in arguments and arguments["pathToSettingsFile"] == None): 

200 

201 if projectName != False: 

202 os.getenv("HOME") 

203 projectDir = os.getenv( 

204 "HOME") + "/.config/%(projectName)s" % locals() 

205 exists = os.path.exists(projectDir) 

206 if not exists: 

207 # Recursively create missing directories 

208 if not os.path.exists(projectDir): 

209 os.makedirs(projectDir) 

210 settingsFile = os.getenv( 

211 "HOME") + "/.config/%(projectName)s/%(projectName)s.yaml" % locals() 

212 exists = os.path.exists(settingsFile) 

213 arguments["settingsFile"] = settingsFile 

214 

215 if not exists: 

216 import codecs 

217 writeFile = codecs.open( 

218 settingsFile, encoding='utf-8', mode='w') 

219 

220 import yaml 

221 # GET CONTENT OF YAML FILE AND REPLACE ~ WITH HOME DIRECTORY 

222 # PATH 

223 with open(settingsFile) as f: 

224 content = f.read() 

225 home = expanduser("~") 

226 content = content.replace("~/", home + "/") 

227 astream = StringIO(content) 

228 

229 if orderedSettings: 

230 this = ordered_load(astream, yaml.SafeLoader) 

231 else: 

232 this = yaml.load(astream) 

233 if this: 

234 

235 settings = this 

236 arguments["<settingsFile>"] = settingsFile 

237 else: 

238 import inspect 

239 ds = os.getcwd() + "/rubbish.yaml" 

240 level = 0 

241 exists = False 

242 count = 1 

243 while not exists and len(ds) and count < 10: 

244 count += 1 

245 level -= 1 

246 exists = os.path.exists(ds) 

247 if not exists: 

248 ds = "/".join(inspect.stack() 

249 [1][1].split("/")[:level]) + "/default_settings.yaml" 

250 

251 shutil.copyfile(ds, settingsFile) 

252 try: 

253 shutil.copyfile(ds, settingsFile) 

254 import codecs 

255 pathToReadFile = settingsFile 

256 try: 

257 readFile = codecs.open( 

258 pathToReadFile, encoding='utf-8', mode='r') 

259 thisData = readFile.read() 

260 readFile.close() 

261 except IOError as e: 

262 message = 'could not open the file %s' % ( 

263 pathToReadFile,) 

264 raise IOError(message) 

265 thisData = thisData.replace( 

266 "/Users/Dave", os.getenv("HOME")) 

267 

268 pathToWriteFile = pathToReadFile 

269 try: 

270 writeFile = codecs.open( 

271 pathToWriteFile, encoding='utf-8', mode='w') 

272 except IOError as e: 

273 message = 'could not open the file %s' % ( 

274 pathToWriteFile,) 

275 raise IOError(message) 

276 writeFile.write(thisData) 

277 writeFile.close() 

278 print( 

279 "default settings have been added to '%(settingsFile)s'. Tailor these settings before proceeding to run %(projectName)s" % locals()) 

280 try: 

281 cmd = """open %(pathToReadFile)s""" % locals() 

282 p = Popen(cmd, stdout=PIPE, 

283 stderr=PIPE, shell=True) 

284 except: 

285 pass 

286 try: 

287 cmd = """start %(pathToReadFile)s""" % locals() 

288 p = Popen(cmd, stdout=PIPE, 

289 stderr=PIPE, shell=True) 

290 except: 

291 pass 

292 except: 

293 print( 

294 "please add settings to file '%(settingsFile)s'" % locals()) 

295 # return 

296 else: 

297 pass 

298 

299 if stream is not False: 

300 import yaml 

301 if orderedSettings: 

302 settings = ordered_load(stream, yaml.SafeLoader) 

303 else: 

304 settings = yaml.load(stream) 

305 

306 # SETUP LOGGER -- DEFAULT TO CONSOLE LOGGER IF NONE PROVIDED IN 

307 # SETTINGS 

308 if 'settings' in locals() and "logging settings" in settings: 

309 if "settingsFile" in arguments: 

310 log = dl.setup_dryx_logging( 

311 yaml_file=arguments["settingsFile"] 

312 ) 

313 elif "<settingsFile>" in arguments: 

314 log = dl.setup_dryx_logging( 

315 yaml_file=arguments["<settingsFile>"] 

316 ) 

317 elif "<pathToSettingsFile>" in arguments: 

318 log = dl.setup_dryx_logging( 

319 yaml_file=arguments["<pathToSettingsFile>"] 

320 ) 

321 elif "--settingsFile" in arguments: 

322 log = dl.setup_dryx_logging( 

323 yaml_file=arguments["--settingsFile"] 

324 ) 

325 elif "pathToSettingsFile" in arguments: 

326 log = dl.setup_dryx_logging( 

327 yaml_file=arguments["pathToSettingsFile"] 

328 ) 

329 

330 elif "--settings" in arguments: 

331 log = dl.setup_dryx_logging( 

332 yaml_file=arguments["--settings"] 

333 ) 

334 

335 elif "--logger" not in arguments or arguments["--logger"] is None: 

336 log = dl.console_logger( 

337 level=self.logLevel 

338 ) 

339 

340 self.log = log 

341 

342 # unpack remaining cl arguments using `exec` to setup the variable names 

343 # automatically 

344 for arg, val in list(arguments.items()): 

345 if arg[0] == "-": 

346 varname = arg.replace("-", "") + "Flag" 

347 else: 

348 varname = arg.replace("<", "").replace(">", "") 

349 if varname == "import": 

350 varname = "iimport" 

351 if isinstance(val, str): 

352 val = val.replace("'", "\\'") 

353 exec(varname + " = '%s'" % (val,)) 

354 else: 

355 exec(varname + " = %s" % (val,)) 

356 if arg == "--dbConn": 

357 dbConn = val 

358 

359 # SETUP A DATABASE CONNECTION BASED ON WHAT ARGUMENTS HAVE BEEN PASSED 

360 dbConn = False 

361 tunnel = False 

362 if ("hostFlag" in locals() and "dbNameFlag" in locals() and hostFlag): 

363 # SETUP DB CONNECTION 

364 dbConn = True 

365 host = arguments["--host"] 

366 user = arguments["--user"] 

367 passwd = arguments["--passwd"] 

368 dbName = arguments["--dbName"] 

369 elif 'settings' in locals() and "database settings" in settings and "host" in settings["database settings"]: 

370 host = settings["database settings"]["host"] 

371 user = settings["database settings"]["user"] 

372 passwd = settings["database settings"]["password"] 

373 dbName = settings["database settings"]["db"] 

374 if "tunnel" in settings["database settings"] and settings["database settings"]["tunnel"]: 

375 tunnel = True 

376 dbConn = True 

377 else: 

378 pass 

379 

380 if not 'settings' in locals(): 

381 settings = False 

382 self.settings = settings 

383 

384 if tunnel: 

385 self._setup_tunnel() 

386 self.dbConn = self.remoteDBConn 

387 return None 

388 

389 if dbConn: 

390 import pymysql as ms 

391 dbConn = ms.connect( 

392 host=host, 

393 user=user, 

394 passwd=passwd, 

395 db=dbName, 

396 use_unicode=True, 

397 charset='utf8', 

398 local_infile=1, 

399 client_flag=ms.constants.CLIENT.MULTI_STATEMENTS, 

400 connect_timeout=36000, 

401 max_allowed_packet=51200000 

402 ) 

403 dbConn.autocommit(True) 

404 

405 self.dbConn = dbConn 

406 

407 return None 

408 

409 def setup( 

410 self): 

411 """ 

412 **Summary:** 

413 *setup the attributes and return* 

414 """ 

415 

416 return self.arguments, self.settings, self.log, self.dbConn 

417 

418 def _setup_tunnel( 

419 self): 

420 """ 

421 *setup ssh tunnel if required* 

422 """ 

423 from subprocess import Popen, PIPE, STDOUT 

424 import pymysql as ms 

425 

426 # SETUP TUNNEL IF REQUIRED 

427 if "ssh tunnel" in self.settings: 

428 # TEST TUNNEL DOES NOT ALREADY EXIST 

429 sshPort = self.settings["ssh tunnel"]["port"] 

430 connected = self._checkServer( 

431 self.settings["database settings"]["host"], sshPort) 

432 if connected: 

433 pass 

434 else: 

435 # GRAB TUNNEL SETTINGS FROM SETTINGS FILE 

436 ru = self.settings["ssh tunnel"]["remote user"] 

437 rip = self.settings["ssh tunnel"]["remote ip"] 

438 rh = self.settings["ssh tunnel"]["remote datbase host"] 

439 

440 cmd = "ssh -fnN %(ru)s@%(rip)s -L %(sshPort)s:%(rh)s:3306" % locals() 

441 p = Popen(cmd, shell=True, close_fds=True) 

442 output = p.communicate()[0] 

443 

444 # TEST CONNECTION - QUIT AFTER SO MANY TRIES 

445 connected = False 

446 count = 0 

447 while not connected: 

448 connected = self._checkServer( 

449 self.settings["database settings"]["host"], sshPort) 

450 time.sleep(1) 

451 count += 1 

452 if count == 5: 

453 self.log.error( 

454 'cound not setup tunnel to remote datbase' % locals()) 

455 sys.exit(0) 

456 

457 if "tunnel" in self.settings["database settings"] and self.settings["database settings"]["tunnel"]: 

458 # TEST TUNNEL DOES NOT ALREADY EXIST 

459 sshPort = self.settings["database settings"]["tunnel"]["port"] 

460 connected = self._checkServer( 

461 self.settings["database settings"]["host"], sshPort) 

462 if connected: 

463 pass 

464 else: 

465 # GRAB TUNNEL SETTINGS FROM SETTINGS FILE 

466 ru = self.settings["database settings"][ 

467 "tunnel"]["remote user"] 

468 rip = self.settings["database settings"]["tunnel"]["remote ip"] 

469 rh = self.settings["database settings"][ 

470 "tunnel"]["remote datbase host"] 

471 

472 cmd = "ssh -fnN %(ru)s@%(rip)s -L %(sshPort)s:%(rh)s:3306" % locals() 

473 p = Popen(cmd, shell=True, close_fds=True) 

474 output = p.communicate()[0] 

475 

476 # TEST CONNECTION - QUIT AFTER SO MANY TRIES 

477 connected = False 

478 count = 0 

479 while not connected: 

480 connected = self._checkServer( 

481 self.settings["database settings"]["host"], sshPort) 

482 time.sleep(1) 

483 count += 1 

484 if count == 5: 

485 self.log.error( 

486 'cound not setup tunnel to remote datbase' % locals()) 

487 sys.exit(0) 

488 

489 # SETUP A DATABASE CONNECTION FOR THE remote database 

490 host = self.settings["database settings"]["host"] 

491 user = self.settings["database settings"]["user"] 

492 passwd = self.settings["database settings"]["password"] 

493 dbName = self.settings["database settings"]["db"] 

494 thisConn = ms.connect( 

495 host=host, 

496 user=user, 

497 passwd=passwd, 

498 db=dbName, 

499 port=sshPort, 

500 use_unicode=True, 

501 charset='utf8', 

502 local_infile=1, 

503 client_flag=ms.constants.CLIENT.MULTI_STATEMENTS, 

504 connect_timeout=36000, 

505 max_allowed_packet=51200000 

506 ) 

507 thisConn.autocommit(True) 

508 self.remoteDBConn = thisConn 

509 

510 return None 

511 

512 def _checkServer(self, address, port): 

513 """ 

514 *Check that the TCP Port we've decided to use for tunnelling is available* 

515 """ 

516 # CREATE A TCP SOCKET 

517 import socket 

518 s = socket.socket() 

519 

520 try: 

521 s.connect((address, port)) 

522 return True 

523 except socket.error as e: 

524 self.log.warning( 

525 """Connection to `%(address)s` on port `%(port)s` failed - try again: %(e)s""" % locals()) 

526 return False 

527 

528 return None 

529 

530 # use the tab-trigger below for new method 

531 # xt-class-method 

532 

533 

534################################################################### 

535# PUBLIC FUNCTIONS # 

536################################################################### 

537def ordered_load(stream, Loader=yaml.Loader, object_pairs_hook=OrderedDict): 

538 class OrderedLoader(Loader): 

539 pass 

540 

541 def construct_mapping(loader, node): 

542 loader.flatten_mapping(node) 

543 return object_pairs_hook(loader.construct_pairs(node)) 

544 OrderedLoader.add_constructor( 

545 yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, 

546 construct_mapping) 

547 return yaml.load(stream, OrderedLoader) 

548 

549 

550if __name__ == '__main__': 

551 main()