|
1 # -*- coding: utf-8 -*- |
|
2 # |
|
3 # MyManager - manage database connections (MySQL version) |
|
4 # |
|
5 # Requires: Python 2.6 / 2.7 / 3.2, MySQLdb |
|
6 # |
|
7 # Part of pydbkit |
|
8 # http://hg.devl.cz/pydbkit |
|
9 # |
|
10 # Copyright (c) 2011, 2013 Radek Brich <radek.brich@devl.cz> |
|
11 # |
|
12 # Permission is hereby granted, free of charge, to any person obtaining a copy |
|
13 # of this software and associated documentation files (the "Software"), to deal |
|
14 # in the Software without restriction, including without limitation the rights |
|
15 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
|
16 # copies of the Software, and to permit persons to whom the Software is |
|
17 # furnished to do so, subject to the following conditions: |
|
18 # |
|
19 # The above copyright notice and this permission notice shall be included in |
|
20 # all copies or substantial portions of the Software. |
|
21 # |
|
22 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
|
23 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
|
24 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
|
25 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
|
26 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
|
27 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN |
|
28 # THE SOFTWARE. |
|
29 |
|
30 """MySQL database connection manager |
|
31 |
|
32 MyManager wraps MySQLdb in same manner as PgManager wraps psycopg2. |
|
33 It's fully compatible so it should work as drop-in replacement for PgManager. |
|
34 |
|
35 It adds following features over MySQLdb: |
|
36 |
|
37 * Save and reuse database connection parameters |
|
38 |
|
39 * Connection pooling |
|
40 |
|
41 * Easy query using the with statement |
|
42 |
|
43 * Dictionary rows |
|
44 |
|
45 Example: |
|
46 |
|
47 from pydbkit import mymanager |
|
48 |
|
49 dbm = mymanager.get_instance() |
|
50 dbm.create_conn(host='127.0.0.1', dbname='default') |
|
51 |
|
52 with dbm.cursor() as curs: |
|
53 curs.execute('SELECT now() AS now') |
|
54 row = curs.fetchone_dict() |
|
55 print(row.now) |
|
56 |
|
57 See PgManager docs for more information. |
|
58 |
|
59 """ |
|
60 |
|
61 from contextlib import contextmanager |
|
62 from collections import OrderedDict |
|
63 import logging |
|
64 import threading |
|
65 import multiprocessing |
|
66 |
|
67 import MySQLdb |
|
68 import MySQLdb.cursors |
|
69 |
|
70 from MySQLdb import DatabaseError, IntegrityError, OperationalError |
|
71 |
|
72 from pydbkit.pgmanager import RowDict |
|
73 |
|
74 |
|
75 log_sql = logging.getLogger("mymanager_sql") |
|
76 log_sql.addHandler(logging.NullHandler()) |
|
77 |
|
78 |
|
79 class MyManagerError(Exception): |
|
80 |
|
81 pass |
|
82 |
|
83 |
|
84 class ConnectionInfo: |
|
85 |
|
86 def __init__(self, name, isolation_level=None, |
|
87 init_statement=None, pool_size=1, **kw): |
|
88 self.name = name # connection name is logged with SQL queries |
|
89 self.isolation_level = isolation_level |
|
90 self.init_statement = init_statement |
|
91 self.pool_size = pool_size |
|
92 self.parameters = kw |
|
93 self.adjust_parameters() |
|
94 |
|
95 def adjust_parameters(self): |
|
96 '''Rename Postgres parameters to proper value for MySQL.''' |
|
97 m = {'dbname' : 'db', 'password' : 'passwd'} |
|
98 res = dict() |
|
99 for k, v in list(self.parameters.items()): |
|
100 if k in m: |
|
101 k = m[k] |
|
102 res[k] = v |
|
103 self.parameters = res |
|
104 |
|
105 |
|
106 class Cursor(MySQLdb.cursors.Cursor): |
|
107 |
|
108 def execute(self, query, args=None): |
|
109 try: |
|
110 return super(Cursor, self).execute(query, args) |
|
111 finally: |
|
112 self._log_query(query, args) |
|
113 |
|
114 def callproc(self, procname, args=None): |
|
115 try: |
|
116 return super(Cursor, self).callproc(procname, args) |
|
117 finally: |
|
118 self._log_query(query, args) |
|
119 |
|
120 def row_dict(self, row, lstrip=None): |
|
121 adjustname = lambda a: a |
|
122 if lstrip: |
|
123 adjustname = lambda a: a.lstrip(lstrip) |
|
124 return RowDict(zip([adjustname(desc[0]) for desc in self.description], row)) |
|
125 |
|
126 def fetchone_dict(self, lstrip=None): |
|
127 row = super(Cursor, self).fetchone() |
|
128 if row is None: |
|
129 return None |
|
130 return self.row_dict(row, lstrip) |
|
131 |
|
132 def fetchall_dict(self, lstrip=None): |
|
133 rows = super(Cursor, self).fetchall() |
|
134 return [self.row_dict(row, lstrip) for row in rows] |
|
135 |
|
136 def mogrify(self, query, args): |
|
137 """Get query with substituted args as it will be send to server.""" |
|
138 if isinstance(query, bytes): |
|
139 query = query.decode() |
|
140 if args is not None: |
|
141 db = self._get_db() |
|
142 query = query % db.literal(args) |
|
143 return query |
|
144 |
|
145 def _log_query(self, query, args): |
|
146 name = self.connection.name if hasattr(self.connection, 'name') else '-' |
|
147 query = self.mogrify(query, args) |
|
148 if isinstance(query, bytes): |
|
149 db = self._get_db() |
|
150 charset = db.character_set_name() |
|
151 query = query.decode(charset) |
|
152 log_sql.debug('[%s] %s' % (name, query)) |
|
153 |
|
154 |
|
155 class MyManager: |
|
156 |
|
157 def __init__(self): |
|
158 self.conn_known = {} # available connections |
|
159 self.conn_pool = {} |
|
160 self.lock = threading.Lock() |
|
161 self.pid = multiprocessing.current_process().pid # forking check |
|
162 |
|
163 def __del__(self): |
|
164 for conn in tuple(self.conn_known.keys()): |
|
165 self.destroy_conn(conn) |
|
166 |
|
167 def create_conn(self, name='default', isolation_level=None, **kw): |
|
168 '''Create named connection.''' |
|
169 if name in self.conn_known: |
|
170 raise MyManagerError('Connection name "%s" already registered.' % name) |
|
171 |
|
172 isolation_level = self._normalize_isolation_level(isolation_level) |
|
173 ci = ConnectionInfo(name, isolation_level, **kw) |
|
174 |
|
175 self.conn_known[name] = ci |
|
176 self.conn_pool[name] = [] |
|
177 |
|
178 def close_conn(self, name='default'): |
|
179 '''Close all connections of given name. |
|
180 |
|
181 Connection credentials are still saved. |
|
182 |
|
183 ''' |
|
184 while len(self.conn_pool[name]): |
|
185 conn = self.conn_pool[name].pop() |
|
186 conn.close() |
|
187 |
|
188 def destroy_conn(self, name='default'): |
|
189 '''Destroy connection. |
|
190 |
|
191 Counterpart of create_conn. |
|
192 |
|
193 ''' |
|
194 if not name in self.conn_known: |
|
195 raise MyManagerError('Connection name "%s" not registered.' % name) |
|
196 |
|
197 self.close_conn(name) |
|
198 |
|
199 del self.conn_known[name] |
|
200 del self.conn_pool[name] |
|
201 |
|
202 def get_conn(self, name='default'): |
|
203 '''Get connection of name 'name' from pool.''' |
|
204 self._check_fork() |
|
205 self.lock.acquire() |
|
206 try: |
|
207 if not name in self.conn_known: |
|
208 raise MyManagerError("Connection name '%s' not registered." % name) |
|
209 |
|
210 # connection from pool |
|
211 conn = None |
|
212 while len(self.conn_pool[name]) and conn is None: |
|
213 conn = self.conn_pool[name].pop() |
|
214 try: |
|
215 conn.ping() |
|
216 except MySQLdb.MySQLError: |
|
217 conn.close() |
|
218 conn = None |
|
219 |
|
220 if conn is None: |
|
221 ci = self.conn_known[name] |
|
222 conn = self._connect(ci) |
|
223 finally: |
|
224 self.lock.release() |
|
225 return conn |
|
226 |
|
227 def put_conn(self, conn, name='default'): |
|
228 '''Put connection back to pool. |
|
229 |
|
230 Name must be same as used for get_conn, |
|
231 otherwise things become broken. |
|
232 |
|
233 ''' |
|
234 self.lock.acquire() |
|
235 try: |
|
236 if not name in self.conn_known: |
|
237 raise MyManagerError("Connection name '%s' not registered." % name) |
|
238 |
|
239 if len(self.conn_pool[name]) >= self.conn_known[name].pool_size: |
|
240 conn.close() |
|
241 return |
|
242 |
|
243 # connection returned to the pool must not be in transaction |
|
244 try: |
|
245 conn.rollback() |
|
246 except OperationalError: |
|
247 conn.close() |
|
248 return |
|
249 |
|
250 self.conn_pool[name].append(conn) |
|
251 finally: |
|
252 self.lock.release() |
|
253 |
|
254 @contextmanager |
|
255 def cursor(self, name='default'): |
|
256 '''Cursor context. |
|
257 |
|
258 Uses any connection of name 'name' from pool |
|
259 and returns cursor for that connection. |
|
260 |
|
261 ''' |
|
262 conn = self.get_conn(name) |
|
263 |
|
264 try: |
|
265 curs = conn.cursor() |
|
266 yield curs |
|
267 finally: |
|
268 curs.close() |
|
269 self.put_conn(conn, name) |
|
270 |
|
271 def _connect(self, ci): |
|
272 conn = MySQLdb.connect(cursorclass=Cursor, **ci.parameters) |
|
273 if not ci.isolation_level is None: |
|
274 if ci.isolation_level == 'AUTOCOMMIT': |
|
275 conn.autocommit(True) |
|
276 else: |
|
277 curs = conn.cursor() |
|
278 curs.execute('SET SESSION TRANSACTION ISOLATION LEVEL ' + ci.isolation_level) |
|
279 curs.close() |
|
280 if ci.init_statement: |
|
281 curs = conn.cursor() |
|
282 curs.execute(ci.init_statement) |
|
283 curs.connection.commit() |
|
284 curs.close() |
|
285 return conn |
|
286 |
|
287 def _normalize_isolation_level(self, level): |
|
288 if level is None: |
|
289 return level |
|
290 if type(level) == str: |
|
291 level = level.upper().replace('_', ' ') |
|
292 if level in ( |
|
293 'AUTOCOMMIT', |
|
294 'READ UNCOMMITTED', |
|
295 'READ COMMITTED', |
|
296 'REPEATABLE READ', |
|
297 'SERIALIZABLE'): |
|
298 return level |
|
299 raise MyManagerError('Unknown isolation level name: "%s"', level) |
|
300 |
|
301 def _check_fork(self): |
|
302 '''Check if process was forked (PID has changed). |
|
303 |
|
304 If it was, clean parent's connections. |
|
305 New connections are created for children. |
|
306 Known connection credentials are inherited, but not shared. |
|
307 |
|
308 ''' |
|
309 if self.pid == multiprocessing.current_process().pid: |
|
310 # PID has not changed |
|
311 return |
|
312 |
|
313 # update saved PID |
|
314 self.pid = multiprocessing.current_process().pid |
|
315 # reinitialize lock |
|
316 self.lock = threading.Lock() |
|
317 # clean parent's connections |
|
318 for name in self.conn_pool: |
|
319 self.conn_pool[name] = [] |
|
320 |
|
321 @classmethod |
|
322 def get_instance(cls): |
|
323 if not hasattr(cls, '_instance'): |
|
324 cls._instance = cls() |
|
325 return cls._instance |
|
326 |
|
327 |
|
328 def get_instance(): |
|
329 return MyManager.get_instance() |
|
330 |