sfa/storage/alchemy.py

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