Allow user to force psycopg version to use, as psycopg2 is not always the best choice

debian/changelog

@@ -1,3 +1,9 @@
+pgloader (2.3.0~dev3-1) experimental; urgency=low
+
+  * Add options for forcing psycopg version to use (-1, -2, --psycopg-version)
+
+ -- Dimitri Fontaine <dim@tapoueh.org>  Wed, 27 Feb 2008 12:54:46 +0100
+
 pgloader (2.3.0~dev2-1) experimental; urgency=low
 
   * columns = * is now supported

debian/control

@@ -8,7 +8,7 @@ Homepage: http://pgloader.projects.postgresql.org
 
 Package: pgloader
 Architecture: all
-Depends: python (>=2.4.4), python-psycopg2 | python-psycopg (<< 1.1.21)
+Depends: python (>=2.4.4), python-psycopg2 | python-psycopg (<< 1.1.21) | python-psycopg (= 1.1.21-14)
 Description: loads flat data files into PostgreSQL
  pgloader imports data from flat files and insert it into a database
  table. You have to provide a configuration file containing sections,

pgloader.1.txt

@@ -145,6 +145,19 @@ You can't use both -F and -I at the same time.
     +/usr/share/pgloader/reformat+. See +reformat_path+ option for
     syntax and default value.
 
+-1, --psycopg1::
+
+    Force usage of psycopg version 1.
+
+-2, --psycopg2::
+
+    Force usage of psycopg version 2.
+
+--psycopg-version::
+
+    Force +pgloader+ to use given version of psycopg, either +1+ or
+    +2+.
+
 Section::
 +
 is the name of a configured Section describing some data to load

pgloader.py

@@ -103,6 +103,20 @@ def parse_options():
                       default = None,
                       help    = "PATH where to find reformat python modules")
 
+    parser.add_option("-1", "--psycopg1", action = "store_true",
+                      dest    = "psycopg1",
+                      default = False,
+                      help    = "Force usage of psycopg1")
+
+    parser.add_option("-2", "--psycopg2", action = "store_true",
+                      dest    = "psycopg2",
+                      default = False,
+                      help    = "Force usage of psycopg2")
+
+    parser.add_option("--psycopg-version", dest = "psycopg_version",
+                      default = None,
+                      help    = "Force usage of given version of psycopg")
+
     (opts, args) = parser.parse_args()
 
     if opts.version:
@@ -132,6 +146,25 @@ def parse_options():
               "Error: Can't be verbose and quiet at the same time!"
         sys.exit(1)
 
+    psyco_opts = 0
+    for opt in [opts.psycopg1, opts.psycopg2, opts.psycopg_version]:
+        if opt:
+            psyco_opts += 1
+
+    if psyco_opts > 1:
+        print >>sys.stderr, \
+              "Error: please use only one of the psycopg options"
+        sys.exit(1)
+
+    if opts.psycopg_version is not None:
+        if opts.psycopg_version in ("1", "2"):
+            opts.psycopg_version = int(opts.psycopg_version)
+
+        else:
+            print >>sys.stderr, \
+                  "Error: psycopg_version can only be set to either 1 or 2"
+            sys.exit(1)
+        
     # if debug, then verbose
     if opts.debug:
         opts.verbose = True
@@ -168,6 +201,14 @@ def parse_options():
     elif opts.quiet:
         pgloader.options.CLIENT_MIN_MESSAGES = logging.ERROR
 
+
+    if opts.psycopg1:
+        pgloader.options.PSYCOPG_VERSION = 1
+    elif opts.psycopg2:
+        pgloader.options.PSYCOPG_VERSION = 2
+    else:
+        pgloader.options.PSYCOPG_VERSION = opts.psycopg_version
+
     return opts.config, args
 
 def parse_config(conffile):

pgloader/db.py

@@ -10,15 +10,42 @@ from options import DRY_RUN, PEDANTIC, CLIENT_MIN_MESSAGES
 from options import TRUNCATE, VACUUM
 from options import INPUT_ENCODING, PG_CLIENT_ENCODING, DATESTYLE
 from options import COPY_SEP, FIELD_SEP, CLOB_SEP, NULL, EMPTY_STRING
+from options import PSYCOPG_VERSION
 
 from tools   import PGLoader_Error
 from logger  import log
 
-try:
-    import psycopg2.psycopg1 as psycopg
-except ImportError:
-    log.info('No psycopg2 module found, trying psycopg1')
-    import psycopg
+log.debug('Preferred psycopg version is %d' % PSYCOPG_VERSION)
+
+if PSYCOPG_VERSION is None:
+    # legacy import behavior
+    log.debug('Trying psycopg2 then psycopg')
+    try:
+        import psycopg2.psycopg1 as psycopg
+    except ImportError:
+        log.info('No psycopg2 module found, trying psycopg1')
+
+        try:
+            import psycopg
+        except ImportError, e:
+            log.fatal('No psycopg module found')
+            raise PGLoader_Error, e
+
+elif PSYCOPG_VERSION == 1:
+    try:
+        log.info("Loading psycopg 1")
+        import psycopg
+    except ImportError, e:
+        log.fatal("Can't load version %d of psycopg" % PSYCOPG_VERSION)
+        raise PGLoader_Error, e
+        
+elif PSYCOPG_VERSION == 2:
+    try:
+        log.info("Loading psycopg 2")
+        import psycopg2.psycopg1 as psycopg
+    except ImportError, e:
+        log.fatal("Can't load version %d of psycopg" % PSYCOPG_VERSION)
+        raise PGLoader_Error, e
 
 class db:
     """ a db connexion and utility class """

pgloader/options.py

@@ -2,7 +2,9 @@
 #
 # Some common options, for each module to get them
 
-PGLOADER_VERSION = '2.3.0~dev2'
+PGLOADER_VERSION = '2.3.0~dev3'
+
+PSYCOPG_VERSION = None
 
 INPUT_ENCODING     = None
 PG_CLIENT_ENCODING = 'latin9'