1#!/usr/local/bin/python 

2# encoding: utf-8 

3""" 

4*a database object that can setup up a ssh tunnel (optional) and a database connection* 

5 

6:Author: 

7 David Young 

8""" 

9from builtins import object 

10import sys 

11import os 

12os.environ['TERM'] = 'vt100' 

13import readline 

14import glob 

15import pickle 

16import time 

17from subprocess import Popen, PIPE, STDOUT 

18from docopt import docopt 

19from fundamentals.mysql import readquery 

20 

21 

22class database(object): 

23 """ 

24 *a database object that can setup up a ssh tunnel (optional) and a database connection* 

25 

26 **Key Arguments** 

27 

28 - ``log`` -- logger 

29 - ``dbSettings`` -- a dictionary of database settings 

30 

31 

32 **Return** 

33 

34 - ``dbConns`` -- a database connection 

35 

36 

37 **Usage** 

38 

39 Given a python dictionary that looks like this: 

40 

41 ```python 

42 dbSettings = { 

43 'host': '127.0.0.1',  

44 'loginPath': 'atlasMovers',  

45 'user': 'monster',  

46 'tunnel': { 

47 'remote ip': 'psweb.mp.qub.ac.uk',  

48 'remote datbase host': 'dormammu',  

49 'remote user': 'monster',  

50 'port': 9006 

51 },  

52 'password': 'myPass',  

53 'db': 'atlas_moving_objects' 

54 } 

55 ``` 

56 

57 ``loginPath`` and ``tunnel`` are optional, to setup the a database connection, run the following: 

58 

59 ```python 

60 # SETUP ALL DATABASE CONNECTIONS 

61 from fundamentals.mysql import database 

62 dbConn = database( 

63 log=log, 

64 dbSettings=dbSettings 

65 ).connect() 

66 ``` 

67 

68 """ 

69 

70 # INITIALISATION 

71 

72 def __init__( 

73 self, 

74 log, 

75 dbSettings=False, 

76 autocommit=True 

77 

78 ): 

79 self.log = log 

80 log.debug("instansiating a new '_database' object") 

81 self.dbSettings = dbSettings 

82 self.autocommit = autocommit 

83 

84 return None 

85 

86 def connect(self): 

87 """*Connect to the database* 

88 

89 **Return** 

90 

91 - ``dbConn`` -- the database connection 

92 

93 

94 See the class docstring for usage 

95 """ 

96 

97 self.log.debug('starting the ``connect`` method') 

98 

99 import pymysql as ms 

100 dbSettings = self.dbSettings 

101 

102 port = False 

103 if "tunnel" in dbSettings and dbSettings["tunnel"]: 

104 port = self._setup_tunnel( 

105 tunnelParameters=dbSettings["tunnel"] 

106 ) 

107 elif "port" in dbSettings and dbSettings["port"]: 

108 port = int(dbSettings["port"]) 

109 

110 # SETUP A DATABASE CONNECTION 

111 host = dbSettings["host"] 

112 user = dbSettings["user"] 

113 passwd = dbSettings["password"] 

114 dbName = dbSettings["db"] 

115 dbConn = ms.connect( 

116 host=host, 

117 user=user, 

118 passwd=passwd, 

119 db=dbName, 

120 port=port, 

121 use_unicode=True, 

122 charset='utf8mb4', 

123 local_infile=1, 

124 client_flag=ms.constants.CLIENT.MULTI_STATEMENTS, 

125 connect_timeout=36000, 

126 max_allowed_packet=51200000 

127 ) 

128 if self.autocommit: 

129 dbConn.autocommit(True) 

130 

131 self.log.debug('completed the ``connect`` method') 

132 return dbConn 

133 

134 def _setup_tunnel( 

135 self, 

136 tunnelParameters): 

137 """ 

138 *Setup a ssh tunnel for a database connection to port through* 

139 

140 **Key Arguments** 

141 

142 - ``tunnelParameters`` -- the tunnel parameters found associated with the database settings 

143 

144 

145 **Return** 

146 

147 - ``sshPort`` -- the port the ssh tunnel is connected via 

148 

149 """ 

150 self.log.debug('starting the ``_setup_tunnel`` method') 

151 

152 # TEST TUNNEL DOES NOT ALREADY EXIST 

153 sshPort = tunnelParameters["port"] 

154 connected = self._checkServer( 

155 "127.0.0.1", sshPort) 

156 if connected: 

157 self.log.debug('ssh tunnel already exists - moving on') 

158 else: 

159 # GRAB TUNNEL SETTINGS FROM SETTINGS FILE 

160 ru = tunnelParameters["remote user"] 

161 rip = tunnelParameters["remote ip"] 

162 rh = tunnelParameters["remote datbase host"] 

163 

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

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

166 output = p.communicate()[0] 

167 self.log.debug('output: %(output)s' % locals()) 

168 

169 # TEST CONNECTION - QUIT AFTER SO MANY TRIES 

170 connected = False 

171 count = 0 

172 while not connected: 

173 connected = self._checkServer( 

174 "127.0.0.1", sshPort) 

175 time.sleep(1) 

176 count += 1 

177 if count == 5: 

178 self.log.error( 

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

180 sys.exit(0) 

181 return sshPort 

182 

183 def _checkServer(self, address, port): 

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

185 """ 

186 self.log.debug('starting the ``_checkServer`` method') 

187 

188 # CREATE A TCP SOCKET 

189 import socket 

190 s = socket.socket() 

191 self.log.debug( 

192 """Attempting to connect to `%(address)s` on port `%(port)s`""" % locals()) 

193 try: 

194 s.connect((address, port)) 

195 self.log.debug( 

196 """Connected to `%(address)s` on port `%(port)s`""" % locals()) 

197 return True 

198 except socket.error as e: 

199 self.log.warning( 

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

201 return False 

202 

203 return None 

204 

205 # xt-class-method