]>
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 | ||
9068dba1 | 11 | import logging |
9068dba1 | 12 | import psycopg2 |
9068dba1 MT |
13 | |
14 | class Connection(object): | |
15 | """ | |
16 | A lightweight wrapper around MySQLdb DB-API connections. | |
17 | ||
18 | The main value we provide is wrapping rows in a dict/object so that | |
19 | columns can be accessed by name. Typical usage:: | |
20 | ||
21 | db = torndb.Connection("localhost", "mydatabase") | |
22 | for article in db.query("SELECT * FROM articles"): | |
23 | print article.title | |
24 | ||
25 | Cursors are hidden by the implementation, but other than that, the methods | |
26 | are very similar to the DB-API. | |
27 | ||
28 | We explicitly set the timezone to UTC and the character encoding to | |
29 | UTF-8 on all connections to avoid time zone and encoding errors. | |
30 | """ | |
31 | def __init__(self, host, database, user=None, password=None): | |
32 | self.host = host | |
33 | self.database = database | |
34 | ||
35 | self._db = None | |
36 | self._db_args = { | |
37 | "host" : host, | |
38 | "database" : database, | |
39 | "user" : user, | |
40 | "password" : password, | |
41 | } | |
42 | ||
43 | try: | |
44 | self.reconnect() | |
45 | except Exception: | |
46 | logging.error("Cannot connect to database on %s", self.host, exc_info=True) | |
47 | ||
48 | def __del__(self): | |
49 | self.close() | |
50 | ||
51 | def close(self): | |
52 | """ | |
53 | Closes this database connection. | |
54 | """ | |
55 | if getattr(self, "_db", None) is not None: | |
56 | self._db.close() | |
57 | self._db = None | |
58 | ||
59 | def reconnect(self): | |
60 | """ | |
61 | Closes the existing database connection and re-opens it. | |
62 | """ | |
63 | self.close() | |
64 | ||
65 | self._db = psycopg2.connect(**self._db_args) | |
66 | self._db.autocommit = True | |
67 | ||
ea324f48 MT |
68 | # Initialize the timezone setting. |
69 | self.execute("SET TIMEZONE TO 'UTC'") | |
70 | ||
9068dba1 MT |
71 | def query(self, query, *parameters, **kwparameters): |
72 | """ | |
73 | Returns a row list for the given query and parameters. | |
74 | """ | |
75 | cursor = self._cursor() | |
76 | try: | |
77 | self._execute(cursor, query, parameters, kwparameters) | |
78 | column_names = [d[0] for d in cursor.description] | |
11347e46 | 79 | return [Row(zip(column_names, row)) for row in cursor] |
9068dba1 MT |
80 | finally: |
81 | cursor.close() | |
82 | ||
83 | def get(self, query, *parameters, **kwparameters): | |
84 | """ | |
85 | Returns the first row returned for the given query. | |
86 | """ | |
87 | rows = self.query(query, *parameters, **kwparameters) | |
88 | if not rows: | |
89 | return None | |
90 | elif len(rows) > 1: | |
91 | raise Exception("Multiple rows returned for Database.get() query") | |
92 | else: | |
93 | return rows[0] | |
94 | ||
95 | def execute(self, query, *parameters, **kwparameters): | |
96 | """ | |
97 | Executes the given query, returning the lastrowid from the query. | |
98 | """ | |
99 | return self.execute_lastrowid(query, *parameters, **kwparameters) | |
100 | ||
101 | def execute_lastrowid(self, query, *parameters, **kwparameters): | |
102 | """ | |
103 | Executes the given query, returning the lastrowid from the query. | |
104 | """ | |
105 | cursor = self._cursor() | |
106 | try: | |
107 | self._execute(cursor, query, parameters, kwparameters) | |
108 | return cursor.lastrowid | |
109 | finally: | |
110 | cursor.close() | |
111 | ||
112 | def execute_rowcount(self, query, *parameters, **kwparameters): | |
113 | """ | |
114 | Executes the given query, returning the rowcount from the query. | |
115 | """ | |
116 | cursor = self._cursor() | |
117 | try: | |
118 | self._execute(cursor, query, parameters, kwparameters) | |
119 | return cursor.rowcount | |
120 | finally: | |
121 | cursor.close() | |
122 | ||
123 | def executemany(self, query, parameters): | |
124 | """ | |
125 | Executes the given query against all the given param sequences. | |
126 | ||
127 | We return the lastrowid from the query. | |
128 | """ | |
129 | return self.executemany_lastrowid(query, parameters) | |
130 | ||
131 | def executemany_lastrowid(self, query, parameters): | |
132 | """ | |
133 | Executes the given query against all the given param sequences. | |
134 | ||
135 | We return the lastrowid from the query. | |
136 | """ | |
137 | cursor = self._cursor() | |
138 | try: | |
139 | cursor.executemany(query, parameters) | |
140 | return cursor.lastrowid | |
141 | finally: | |
142 | cursor.close() | |
143 | ||
144 | def executemany_rowcount(self, query, parameters): | |
145 | """ | |
146 | Executes the given query against all the given param sequences. | |
147 | ||
148 | We return the rowcount from the query. | |
149 | """ | |
150 | cursor = self._cursor() | |
151 | ||
152 | try: | |
153 | cursor.executemany(query, parameters) | |
154 | return cursor.rowcount | |
155 | finally: | |
156 | cursor.close() | |
157 | ||
158 | def _ensure_connected(self): | |
159 | if self._db is None: | |
a69e87a1 | 160 | logging.warning("Database connection was lost...") |
d06b2491 | 161 | |
9068dba1 MT |
162 | self.reconnect() |
163 | ||
164 | def _cursor(self): | |
165 | self._ensure_connected() | |
166 | return self._db.cursor() | |
167 | ||
168 | def _execute(self, cursor, query, parameters, kwparameters): | |
053aeef1 | 169 | logging.debug("SQL Query: %s" % (query % (kwparameters or parameters))) |
66862195 | 170 | |
9068dba1 MT |
171 | try: |
172 | return cursor.execute(query, kwparameters or parameters) | |
747e47db | 173 | except (OperationalError, psycopg2.ProgrammingError): |
9068dba1 MT |
174 | logging.error("Error connecting to database on %s", self.host) |
175 | self.close() | |
176 | raise | |
177 | ||
a69e87a1 MT |
178 | def transaction(self): |
179 | return Transaction(self) | |
280ed84e | 180 | |
9068dba1 MT |
181 | |
182 | class Row(dict): | |
183 | """A dict that allows for object-like property access syntax.""" | |
184 | def __getattr__(self, name): | |
185 | try: | |
186 | return self[name] | |
187 | except KeyError: | |
188 | raise AttributeError(name) | |
189 | ||
190 | ||
280ed84e | 191 | class Transaction(object): |
a69e87a1 MT |
192 | def __init__(self, db): |
193 | self.db = db | |
280ed84e | 194 | |
a69e87a1 | 195 | self.db.execute("START TRANSACTION") |
280ed84e | 196 | |
a69e87a1 MT |
197 | def __enter__(self): |
198 | return self | |
280ed84e | 199 | |
a69e87a1 MT |
200 | def __exit__(self, exctype, excvalue, traceback): |
201 | if exctype is not None: | |
202 | self.db.execute("ROLLBACK") | |
203 | else: | |
204 | self.db.execute("COMMIT") | |
280ed84e MT |
205 | |
206 | ||
9068dba1 MT |
207 | # Alias some common exceptions |
208 | IntegrityError = psycopg2.IntegrityError | |
209 | OperationalError = psycopg2.OperationalError |