sfa/storage/alchemy.py

   1 from sqlalchemy import create_engine
   2 from sqlalchemy.orm import sessionmaker
   3
   4 from sqlalchemy.ext.declarative import declarative_base
   5 from sqlalchemy import Column, Integer, String
   6 from sqlalchemy.orm import relationship, backref
   7 from sqlalchemy import ForeignKey
   8
   9 from sfa.util.sfalogging import logger
  10
  11 Base=declarative_base()
  12
  13 class Alchemy:
  14
  15     def __init__ (self, config):
  16         dbname="sfa"
  17         # will be created lazily on-demand
  18         self._session = None
  19         # the former PostgreSQL.py used the psycopg2 directly and was doing
  20         #self.connection.set_client_encoding("UNICODE")
  21         # it's unclear how to achieve this in sqlalchemy, nor if it's needed at all
  22         # http://www.sqlalchemy.org/docs/dialects/postgresql.html#unicode
  23         # we indeed have /var/lib/pgsql/data/postgresql.conf where
  24         # this setting is unset, it might be an angle to tweak that if need be
  25         # try a unix socket first - omitting the hostname does the trick
  26         unix_desc = "postgresql+psycopg2://%s:%s@:%s/%s"%\
  27             (config.SFA_DB_USER,config.SFA_DB_PASSWORD,config.SFA_DB_PORT,dbname)
  28         # the TCP fallback method
  29         tcp_desc = "postgresql+psycopg2://%s:%s@%s:%s/%s"%\
  30             (config.SFA_DB_USER,config.SFA_DB_PASSWORD,config.SFA_DB_HOST,config.SFA_DB_PORT,dbname)
  31         for engine_desc in [ unix_desc, tcp_desc ] :
  32             try:
  33                 self.engine = create_engine (engine_desc)
  34                 self.check()
  35                 return
  36             except:
  37                 pass
  38         self.engine=None
  39         raise Exception,"Could not connect to database"
  40
  41
  42     # expects boolean True: debug is ON or False: debug is OFF
  43     def debug (self, echo):
  44         self.engine.echo=echo
  45
  46     def check (self):
  47         self.engine.execute ("select 1").scalar()
  48
  49     # create schema
  50     # warning: need to have all Base subclass loaded for this to work
  51     def create_schema (self):
  52         return Base.metadata.create_all(self.engine)
  53
  54     # does a complete wipe of the schema, use with care
  55     def drop_schema (self):
  56         return Base.metadata.drop_all(self.engine)
  57
  58     def session (self):
  59         if self._session is None:
  60             Session=sessionmaker ()
  61             self._session=Session(bind=self.engine)
  62         return self._session
  63
  64     def close_session (self):
  65         if self._session is None: return
  66         self._session.close()
  67         self._session=None
  68
  69     def commit (self):
  70         self.session().commit()
  71
  72     def insert (self, stuff, commit=False):
  73         if isinstance (stuff,list):
  74             self.session().add_all(stuff)
  75         else:
  76             self.session().add(obj)
  77
  78     # for compat with the previous PostgreSQL stuff
  79     def update (self, record):
  80         self.commit()
  81
  82     def remove (self, record):
  83         del record
  84         self.commit()
  85
  86 ####################
  87 # dicts vs objects
  88 ####################
  89 # historically the front end to the db dealt with dicts, so the code was only dealing with dicts
  90 # sqlalchemy however offers an object interface, meaning that you write obj.id instead of obj['id']
  91 # which is admittedly much nicer
  92 # however we still need to deal with dictionaries if only for the xmlrpc layer
  93 #
  94 # here are a few utilities for this
  95 #
  96 # (*) first off, when an old pieve of code needs to be used as-is, if only temporarily, the simplest trick
  97 # is to use obj.__dict__
  98 # this behaves exactly like required, i.e. obj.__dict__['field']='new value' does change obj.field
  99 # however this depends on sqlalchemy's implementation so it should be avoided
 100 #
 101 # (*) second, when an object needs to be exposed to the xmlrpc layer, we need to convert it into a dict
 102 # remember though that writing the resulting dictionary won't change the object
 103 # essentially obj.__dict__ would be fine too, except that we want to discard alchemy private keys starting with '_'
 104 # 2 ways are provided for that:
 105 # . dict(obj)
 106 # . obj.todict()
 107 # the former dict(obj) relies on __iter__() and next() below, and does not rely on the fields names
 108 # although it seems to work fine, I've found cases where it issues a weird python error that I could not get right
 109 # so the latter obj.todict() seems more reliable but more hacky as is relies on the form of fields, so this can probably be improved
 110 #
 111 # (*) finally for converting a dictionary into an sqlalchemy object, we provide
 112 # obj.set_from_dict(dict)
 113
 114 from sqlalchemy.orm import object_mapper
 115 class AlchemyObj:
 116     def __iter__(self):
 117         self._i = iter(object_mapper(self).columns)
 118         return self
 119     def next(self):
 120         n = self._i.next().name
 121         return n, getattr(self, n)
 122     def todict (self):
 123         d=self.__dict__
 124         keys=[k for k in d.keys() if not k.startswith('_')]
 125         return dict ( [ (k,d[k]) for k in keys ] )
 126     def set_from_dict (self, d):
 127         for (k,v) in d.iteritems():
 128             setattr(self,k,v)
 129
 130 ####################
 131 from sfa.util.config import Config
 132
 133 alchemy=Alchemy (Config())
 134 engine=alchemy.engine
 135 dbsession=alchemy.session()
 136