]>
Commit | Line | Data |
---|---|---|
9068dba1 MT |
1 | #!/usr/bin/env python |
2 | ||
3 | """ | |
4 | A lightweight wrapper around psycopg2. | |
5 | ||
6 | Originally part of the Tornado framework. The tornado.database module | |
7 | is slated for removal in Tornado 3.0, and it is now available separately | |
8 | as torndb. | |
9 | """ | |
10 | ||
11 | from __future__ import absolute_import, division, with_statement | |
12 | ||
13 | import copy | |
14 | import itertools | |
15 | import logging | |
16 | import os | |
17 | import psycopg2 | |
18 | import time | |
19 | ||
20 | class Connection(object): | |
21 | """ | |
22 | A lightweight wrapper around MySQLdb DB-API connections. | |
23 | ||
24 | The main value we provide is wrapping rows in a dict/object so that | |
25 | columns can be accessed by name. Typical usage:: | |
26 | ||
27 | db = torndb.Connection("localhost", "mydatabase") | |
28 | for article in db.query("SELECT * FROM articles"): | |
29 | print article.title | |
30 | ||
31 | Cursors are hidden by the implementation, but other than that, the methods | |
32 | are very similar to the DB-API. | |
33 | ||
34 | We explicitly set the timezone to UTC and the character encoding to | |
35 | UTF-8 on all connections to avoid time zone and encoding errors. | |
36 | """ | |
37 | def __init__(self, host, database, user=None, password=None): | |
38 | self.host = host | |
39 | self.database = database | |
40 | ||
41 | self._db = None | |
42 | self._db_args = { | |
43 | "host" : host, | |
44 | "database" : database, | |
45 | "user" : user, | |
46 | "password" : password, | |
47 | } | |
48 | ||
49 | try: | |
50 | self.reconnect() | |
51 | except Exception: | |
52 | logging.error("Cannot connect to database on %s", self.host, exc_info=True) | |
53 | ||
54 | def __del__(self): | |
55 | self.close() | |
56 | ||
57 | def close(self): | |
58 | """ | |
59 | Closes this database connection. | |
60 | """ | |
61 | if getattr(self, "_db", None) is not None: | |
62 | self._db.close() | |
63 | self._db = None | |
64 | ||
65 | def reconnect(self): | |
66 | """ | |
67 | Closes the existing database connection and re-opens it. | |
68 | """ | |
69 | self.close() | |
70 | ||
71 | self._db = psycopg2.connect(**self._db_args) | |
72 | self._db.autocommit = True | |
73 | ||
74 | def query(self, query, *parameters, **kwparameters): | |
75 | """ | |
76 | Returns a row list for the given query and parameters. | |
77 | """ | |
78 | cursor = self._cursor() | |
79 | try: | |
80 | self._execute(cursor, query, parameters, kwparameters) | |
81 | column_names = [d[0] for d in cursor.description] | |
82 | return [Row(itertools.izip(column_names, row)) for row in cursor] | |
83 | finally: | |
84 | cursor.close() | |
85 | ||
86 | def get(self, query, *parameters, **kwparameters): | |
87 | """ | |
88 | Returns the first row returned for the given query. | |
89 | """ | |
90 | rows = self.query(query, *parameters, **kwparameters) | |
91 | if not rows: | |
92 | return None | |
93 | elif len(rows) > 1: | |
94 | raise Exception("Multiple rows returned for Database.get() query") | |
95 | else: | |
96 | return rows[0] | |
97 | ||
98 | def execute(self, query, *parameters, **kwparameters): | |
99 | """ | |
100 | Executes the given query, returning the lastrowid from the query. | |
101 | """ | |
102 | return self.execute_lastrowid(query, *parameters, **kwparameters) | |
103 | ||
104 | def execute_lastrowid(self, query, *parameters, **kwparameters): | |
105 | """ | |
106 | Executes the given query, returning the lastrowid from the query. | |
107 | """ | |
108 | cursor = self._cursor() | |
109 | try: | |
110 | self._execute(cursor, query, parameters, kwparameters) | |
111 | return cursor.lastrowid | |
112 | finally: | |
113 | cursor.close() | |
114 | ||
115 | def execute_rowcount(self, query, *parameters, **kwparameters): | |
116 | """ | |
117 | Executes the given query, returning the rowcount from the query. | |
118 | """ | |
119 | cursor = self._cursor() | |
120 | try: | |
121 | self._execute(cursor, query, parameters, kwparameters) | |
122 | return cursor.rowcount | |
123 | finally: | |
124 | cursor.close() | |
125 | ||
126 | def executemany(self, query, parameters): | |
127 | """ | |
128 | Executes the given query against all the given param sequences. | |
129 | ||
130 | We return the lastrowid from the query. | |
131 | """ | |
132 | return self.executemany_lastrowid(query, parameters) | |
133 | ||
134 | def executemany_lastrowid(self, query, parameters): | |
135 | """ | |
136 | Executes the given query against all the given param sequences. | |
137 | ||
138 | We return the lastrowid from the query. | |
139 | """ | |
140 | cursor = self._cursor() | |
141 | try: | |
142 | cursor.executemany(query, parameters) | |
143 | return cursor.lastrowid | |
144 | finally: | |
145 | cursor.close() | |
146 | ||
147 | def executemany_rowcount(self, query, parameters): | |
148 | """ | |
149 | Executes the given query against all the given param sequences. | |
150 | ||
151 | We return the rowcount from the query. | |
152 | """ | |
153 | cursor = self._cursor() | |
154 | ||
155 | try: | |
156 | cursor.executemany(query, parameters) | |
157 | return cursor.rowcount | |
158 | finally: | |
159 | cursor.close() | |
160 | ||
161 | def _ensure_connected(self): | |
162 | if self._db is None: | |
163 | self.reconnect() | |
164 | ||
165 | def _cursor(self): | |
166 | self._ensure_connected() | |
167 | return self._db.cursor() | |
168 | ||
169 | def _execute(self, cursor, query, parameters, kwparameters): | |
170 | try: | |
171 | return cursor.execute(query, kwparameters or parameters) | |
172 | except OperationalError: | |
173 | logging.error("Error connecting to database on %s", self.host) | |
174 | self.close() | |
175 | raise | |
176 | ||
177 | ||
178 | class Row(dict): | |
179 | """A dict that allows for object-like property access syntax.""" | |
180 | def __getattr__(self, name): | |
181 | try: | |
182 | return self[name] | |
183 | except KeyError: | |
184 | raise AttributeError(name) | |
185 | ||
186 | ||
187 | # Alias some common exceptions | |
188 | IntegrityError = psycopg2.IntegrityError | |
189 | OperationalError = psycopg2.OperationalError |