1 # -*- coding: utf-8 -*- |
1 # -*- coding: utf-8 -*- |
2 # |
2 # |
3 # PgManager - manage database connections |
3 # PgManager - manage database connections |
4 # |
4 # |
5 # Requires: Python 2.6, psycopg2 |
5 # Requires: Python 3.2, psycopg2 |
6 # |
6 # |
7 # Part of pgtoolkit |
7 # Part of pgtoolkit |
8 # http://hg.devl.cz/pgtoolkit |
8 # http://hg.devl.cz/pgtoolkit |
9 # |
9 # |
10 # Copyright (c) 2010, 2011 Radek Brich <radek.brich@devl.cz> |
10 # Copyright (c) 2010, 2011 Radek Brich <radek.brich@devl.cz> |
68 The row returned by fetchone_dict() is special dict object, which can be accessed |
68 The row returned by fetchone_dict() is special dict object, which can be accessed |
69 using item or attribute access, that is row['now'] or row.now. |
69 using item or attribute access, that is row['now'] or row.now. |
70 """ |
70 """ |
71 |
71 |
72 from contextlib import contextmanager |
72 from contextlib import contextmanager |
|
73 from collections import OrderedDict |
73 import logging |
74 import logging |
74 import threading |
75 import threading |
75 import multiprocessing |
76 import multiprocessing |
76 import select |
77 import select |
77 import socket |
78 import socket |
101 self.keep_alive = keep_alive |
102 self.keep_alive = keep_alive |
102 self.init_statement = init_statement |
103 self.init_statement = init_statement |
103 self.keep_open = keep_open |
104 self.keep_open = keep_open |
104 |
105 |
105 |
106 |
106 class RowDict(dict): |
107 class RowDict(OrderedDict): |
107 |
108 |
108 def __getattr__(self, key): |
109 def __getattr__(self, key): |
109 return self[key] |
110 try: |
|
111 return self[key] |
|
112 except KeyError: |
|
113 raise AttributeError(key) |
110 |
114 |
111 |
115 |
112 class Cursor(psycopg2.extensions.cursor): |
116 class Cursor(psycopg2.extensions.cursor): |
113 |
117 |
114 def execute(self, query, args=None): |
118 def execute(self, query, args=None): |
130 if lstrip: |
134 if lstrip: |
131 adjustname = lambda a: a.lstrip(lstrip) |
135 adjustname = lambda a: a.lstrip(lstrip) |
132 return RowDict(zip([adjustname(desc[0]) for desc in self.description], row)) |
136 return RowDict(zip([adjustname(desc[0]) for desc in self.description], row)) |
133 |
137 |
134 def fetchone_dict(self, lstrip=None): |
138 def fetchone_dict(self, lstrip=None): |
|
139 '''Return one row as OrderedDict''' |
135 row = super(Cursor, self).fetchone() |
140 row = super(Cursor, self).fetchone() |
136 if row is None: |
141 if row is None: |
137 return None |
142 return None |
138 return self.row_dict(row, lstrip) |
143 return self.row_dict(row, lstrip) |
139 |
144 |
140 def fetchall_dict(self, lstrip=None): |
145 def fetchall_dict(self, lstrip=None): |
|
146 '''Return all rows as OrderedDict''' |
141 rows = super(Cursor, self).fetchall() |
147 rows = super(Cursor, self).fetchall() |
142 return [self.row_dict(row, lstrip) for row in rows] |
148 return [self.row_dict(row, lstrip) for row in rows] |
143 |
149 |
144 def fetchone_adapted(self): |
150 def fetchone_adapted(self, lstrip=None): |
145 '''Like fetchone() but values are quoted for direct inclusion in SQL query. |
151 '''Like fetchone_dict() but values are quoted for direct inclusion in SQL query. |
146 |
152 |
147 This is useful when you need to generate SQL script from data returned |
153 This is useful when you need to generate SQL script from data returned |
148 by the query. Use mogrify() for simple cases. |
154 by the query. Use mogrify() for simple cases. |
149 |
155 |
150 ''' |
156 ''' |
151 row = super(Cursor, self).fetchone() |
157 row = super(Cursor, self).fetchone() |
152 if row is None: |
158 if row is None: |
153 return None |
159 return None |
154 return [self.mogrify('%s', [x]).decode('utf8') for x in row] |
160 return self.row_dict([self.mogrify('%s', [x]).decode('utf8') for x in row], lstrip) |
155 |
161 |
156 def fetchall_adapted(self): |
162 def fetchall_adapted(self, lstrip=None): |
157 '''Like fetchall() but values are quoted for direct inclusion in SQL query.''' |
163 '''Like fetchall_dict() but values are quoted for direct inclusion in SQL query.''' |
158 rows = super(Cursor, self).fetchall() |
164 rows = super(Cursor, self).fetchall() |
159 return [[self.mogrify('%s', [x]).decode('utf8') for x in row] for row in rows] |
165 return [self.row_dict([self.mogrify('%s', [x]).decode('utf8') for x in row], lstrip) for row in rows] |
160 |
166 |
161 |
167 |
162 class Connection(psycopg2.extensions.connection): |
168 class Connection(psycopg2.extensions.connection): |
163 |
169 |
164 def cursor(self, name=None): |
170 def cursor(self, name=None): |