dbwrap.py

   1 class bag(dict):
   2         """A simple dict-derived class that allows bag.key and bag['key']"""
   3         def __getattr__(self, key):
   4                 return self[key]
   5
   6         def __setattr__(self, key, value):
   7                 self[key] = value
   8
   9         def __delattr__(self, key):
  10                 del self[key]
  11
  12 class rows_ref(object):
  13         """A lazy reference to a set of rows.
  14
  15         Again, this should not be instantiated by hand; use table.rows(...)
  16         """
  17         def __init__(self, con, table, where_clause, args):
  18                 self.con, self.table, self.where_clause, self.args = con, table, where_clause, args
  19
  20         def _set_clause(self, **kwargs):
  21                 return ', '.join('%s = %s' % (column, self.con.placeholder) for column in kwargs.keys())
  22
  23         def select(self, *args):
  24                 """Returns all rows in this rows_ref. If arguments are given, they are the columns selected."""
  25                 return self.con.query('SELECT %s FROM %s WHERE %s' % (
  26                                 ', '.join(args) if args else '*',
  27                                 self.table.name,
  28                                 self.where_clause,
  29                         ),
  30                         *self.args
  31                 )
  32
  33         def select_one(self, *args):
  34                 """Returns the first row in this rows_ref. If arguments are given, they are the columns selected."""
  35                 return self.con.query_one('SELECT %s FROM %s WHERE %s' % (
  36                                 ', '.join(args) if args else '*',
  37                                 self.table.name,
  38                                 self.where_clause,
  39                         ),
  40                         *self.args
  41                 )
  42
  43         def select_column(self, column):
  44                 """Returns the values in the specified column as a list."""
  45                 results = self.con.query('SELECT %s FROM %s WHERE %s' % (
  46                                 column,
  47                                 self.table.name,
  48                                 self.where_clause,
  49                         ),
  50                         *self.args
  51                 )
  52
  53                 if not results:
  54                         return []
  55
  56                 return [row[column] for row in result]
  57
  58         def select_value(self, column):
  59                 """Returns the value in the specified column."""
  60                 result = self.con.query_one('SELECT %s FROM %s WHERE %s' % (
  61                                 column,
  62                                 self.table.name,
  63                                 self.where_clause,
  64                         ),
  65                         *self.args
  66                 )
  67
  68                 if not result:
  69                         return None
  70
  71                 return result[column]
  72
  73         def update(self, _clause = '', *args, **values):
  74                 """Updates rows in this rows_ref.
  75
  76                 Note: this uses similar arguments to table.select()"""
  77                 self.con.execute('UPDATE %s SET %s WHERE %s' % (
  78                                 self.table.name,
  79                                 _clause if _clause else self._set_clause(**values),
  80                                 self.where_clause),
  81                         *((args if _clause else tuple(values.values())) + tuple(self.args))
  82                 )
  83
  84         def delete(self):
  85                 """Deletes all rows in this rows_ref"""
  86                 self.con.execute('DELETE FROM %s WHERE %s' % (self.table.name, self.where_clause), *self.args)
  87
  88         def exist(self):
  89                 """Returns whether this is an empty set"""
  90                 return self.table.select_one(self.where_clause, *self.args) != None
  91
  92 class table_ref(object):
  93         """A reference to a table.
  94
  95         Note that you shouldn't instantiate this yourself; use wrapper['table-name'] or wrapper.table_name instead.
  96         """
  97         def __init__(self, con, name):
  98                 self.con, self.name = con, name
  99
 100         def _where_clause(self, *args, **criteria):
 101                 return ' AND '.join(('%s = ' + ('%' + self.con.placeholder if self.con.placeholder.startswith('%') else self.con.placeholder) if value is not None else '%s IS NULL') % column for column, value in zip(criteria.keys(), args or criteria.values()))
 102
 103         def _create_clause(self, _query, args, criteria):
 104                 return _query if _query else self._where_clause(*args, **criteria)
 105
 106         def insert(self, _columns = [], _clause = '', *args, **values):
 107                 """INSERTs a row into the table.
 108
 109                 If columns and _clause are given, they will be placed into the SQL statement; otherwise, the columns and VALUES will be automatically generated.
 110                 Note that columns _must_ be given if clause is.
 111                 """
 112
 113                 self.con.execute('INSERT INTO %s(%s) %s' % (self.name, ', '.join(_columns if _columns else values.keys()), _clause if _clause else 'VALUES(%s)' % ', '.join([self.con.placeholder] * len(args if args else values))), *(values.values() if values else args))
 114
 115         def all(self):
 116                 """Returns the whole table"""
 117
 118                 return self.con.query('SELECT * FROM %s' % self.name)
 119
 120         def select(self, _query = '', *args, **criteria):
 121                 """SELECTs from the table.
 122
 123                 If _query is specified, it will be used as the where clause; otherwise, one will be generated from the keyword arguments.
 124                 """
 125
 126                 return self.rows(_query, *args, **criteria).select()
 127
 128         def select_one(self, _query = '', *args, **criteria):
 129                 """SELECTs from the table.
 130
 131                 If _query is specified, it will be used as the where clause; otherwise, one will be generated from the keyword arguments.
 132                 """
 133                 return self.rows(_query, *args, **criteria).select_one()
 134
 135         def rows(self, _query = '', *args, **criteria):
 136                 """Returns a rows_ref object; uses same arguments as select()."""
 137                 return rows_ref(self.con, self, self._create_clause(_query, args, criteria), (args if _query else tuple(x for x in criteria.values() if x is not None)))
 138
 139 class wrapper(object):
 140         """A wrapper for a DB-API connection. You MUST specify the placeholder for parameters in the call to __init__; this is ? for most libraries and %s for MySQLdb.
 141
 142         NOTE: If your table names are not valid Python identifiers, use wrapper['table-name'].
 143         """
 144         def __init__(self, con, placeholder, debug = False):
 145                 self.con, self.debug = con, debug
 146                 self.placeholder = placeholder
 147
 148         def __getattr__(self, key):
 149                 """Returns a table object"""
 150                 return self.__dict__[key] if key in self.__dict__ else  table_ref(self, key)
 151
 152         def commit(self):
 153                 """Commits the transaction"""
 154                 self.con.commit()
 155
 156         def execute(self, query, *args):
 157                 """Runs self.con.cursor().execute(query, *args), and returns the cursor"""
 158
 159                 if self.debug: print query, args
 160
 161                 cur = self.con.cursor()
 162                 cur.execute(query, args)
 163                 return cur
 164
 165         def execute_script(self, filename):
 166                 """Runs the SQL script in 'filename'"""
 167
 168                 script = file(filename, 'r')
 169
 170                 while True:
 171                         data = ''
 172                         while not data.rstrip().endswith(';'):
 173                                 data += script.readline()
 174                                 if data == '': return
 175
 176                         self.execute(data)
 177
 178         def query(self, query, *args):
 179                 """Wrapper similar to execute for SELECT statements, that returns a list of dicts."""
 180
 181                 cur = None
 182
 183                 if self.debug:
 184                         print query, args
 185                         self.debug = False
 186                         cur = self.execute(query, *args)
 187                         self.debug = True
 188                 else:
 189                         cur = self.execute(query, *args)
 190
 191                 results = cur.fetchall()
 192                 if cur.description and results:
 193                         # all SELECTs should have a description, but we shouldn't choke as the result of idiots
 194                         columns = [desc[0] for desc in cur.description]
 195                         return [bag(zip(columns, result)) for result in results]
 196                 else: return []
 197
 198         def query_one(self, query, *args):
 199                 """Derivative of query for one-row SELECTs."""
 200
 201                 cur = None
 202
 203                 if self.debug:
 204                         print query, args
 205                         self.debug = False
 206                         cur = self.execute(query, *args)
 207                         self.debug = True
 208                 else:
 209                         cur = self.execute(query, *args)
 210
 211                 result = cur.fetchone()
 212                 if cur.description and result:
 213                         columns = (desc[0] for desc in cur.description)
 214                         return bag(zip(columns, result))
 215                 else:
 216                         return bag()
 217
 218         __getitem__ = __getattr__
 219
 220         def __del__(self):
 221                 self.con.close()
 222                 del self.con
 223