X-Git-Url: https://git.decadent.org.uk/gitweb/?a=blobdiff_plain;ds=sidebyside;f=daklib%2Fdbconn.py;h=979256e0c5c0969c24d94fd1c60df57f40baabfb;hb=e6a0eee330810ba5bdb752f67f7dc6cc2b623184;hp=1bf5a0f997c7bb84aa5821aa4f9ab3c074ec3422;hpb=8c93e61298c8d126f9ceb7cdab49ae6354de74f5;p=dak.git

diff --git a/daklib/dbconn.py b/daklib/dbconn.py
index 1bf5a0f9..979256e0 100755
--- a/daklib/dbconn.py
+++ b/daklib/dbconn.py
@@ -59,7 +59,7 @@ import sqlalchemy
 from sqlalchemy import create_engine, Table, MetaData, Column, Integer, desc, \
     Text, ForeignKey
 from sqlalchemy.orm import sessionmaker, mapper, relation, object_session, \
-    backref, MapperExtension, EXT_CONTINUE, object_mapper
+    backref, MapperExtension, EXT_CONTINUE, object_mapper, clear_mappers
 from sqlalchemy import types as sqltypes
 
 # Don't remove this, we re-export the exceptions to scripts which import us
@@ -3164,16 +3164,18 @@ class DBConn(object):
     def __createconn(self):
         from config import Config
         cnf = Config()
-        if cnf["DB::Host"]:
+        if cnf.has_key("DB::Service"):
+            connstr = "postgresql://service=%s" % cnf["DB::Service"]
+        elif cnf.has_key("DB::Host"):
             # TCP/IP
-            connstr = "postgres://%s" % cnf["DB::Host"]
-            if cnf["DB::Port"] and cnf["DB::Port"] != "-1":
+            connstr = "postgresql://%s" % cnf["DB::Host"]
+            if cnf.has_key("DB::Port") and cnf["DB::Port"] != "-1":
                 connstr += ":%s" % cnf["DB::Port"]
             connstr += "/%s" % cnf["DB::Name"]
         else:
             # Unix Socket
-            connstr = "postgres:///%s" % cnf["DB::Name"]
-            if cnf["DB::Port"] and cnf["DB::Port"] != "-1":
+            connstr = "postgresql:///%s" % cnf["DB::Name"]
+            if cnf.has_key("DB::Port") and cnf["DB::Port"] != "-1":
                 connstr += "?port=%s" % cnf["DB::Port"]
 
         engine_args = { 'echo': self.debug }
@@ -3185,6 +3187,20 @@ class DBConn(object):
             cnf['DB::Unicode'] == 'false':
             engine_args['use_native_unicode'] = False
 
+        # Monkey patch a new dialect in in order to support service= syntax
+        import sqlalchemy.dialects.postgresql
+        from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
+        class PGDialect_psycopg2_dak(PGDialect_psycopg2):
+            def create_connect_args(self, url):
+                if str(url).startswith('postgresql://service='):
+                    # Eww
+                    servicename = str(url)[21:]
+                    return (['service=%s' % servicename], {})
+                else:
+                    return PGDialect_psycopg2.create_connect_args(self, url)
+
+        sqlalchemy.dialects.postgresql.base.dialect = PGDialect_psycopg2_dak
+
         self.db_pg   = create_engine(connstr, **engine_args)
         self.db_meta = MetaData()
         self.db_meta.bind = self.db_pg
@@ -3198,6 +3214,15 @@ class DBConn(object):
     def session(self):
         return self.db_smaker()
 
+    def reset(self):
+        '''
+        Resets the DBConn object. This function must be called by subprocesses
+        created by the multiprocessing module. See tests/dbtest_multiproc.py
+        for an example.
+        '''
+        clear_mappers()
+        self.__createconn()
+
 __all__.append('DBConn')