ldap: add option to disable referral following

mylib/ldap.py

@@ -60,12 +60,14 @@ class LdapServer:
     con = 0
 
     def __init__(self, uri, dn=None, pwd=None, v2=None,
-                 raiseOnError=False, logger=False, checkCert=True):
+                 raiseOnError=False, logger=False, checkCert=True,
+                 disableReferral=False):
         self.uri = uri
         self.dn = dn
         self.pwd = pwd
         self.raiseOnError = raiseOnError
         self.checkCert = checkCert
+        self.disableReferral = disableReferral
         if v2:
             self.v2 = True
         if logger:
@@ -85,6 +87,9 @@ class LdapServer:
                 if not self.checkCert:
                     # pylint: disable=no-member
                     ldap.set_option(ldap.OPT_X_TLS_REQUIRE_CERT, ldap.OPT_X_TLS_NEVER)
+                if self.disableReferral:
+                    # pylint: disable=no-member
+                    ldap.set_option(ldap.OPT_REFERRALS, ldap.OPT_OFF)
                 con = ldap.initialize(self.uri)
                 if self.v2:
                     con.protocol_version = ldap.VERSION2  # pylint: disable=no-member
@@ -468,6 +473,9 @@ class LdapClient:
         section.add_option(
             BooleanOption, 'checkcert', default=True,
             comment='Check LDAP certificate')
+        section.add_option(
+            BooleanOption, 'disablereferral', default=False,
+            comment='Disable referral following')
 
         return section
 
@@ -480,7 +488,9 @@ class LdapClient:
         log.info("Connect to LDAP server %s as %s", uri, binddn if binddn else 'annonymous')
         self._conn = LdapServer(
             uri, dn=binddn, pwd=self._get_option('bindpwd'),
-            checkCert=self._get_option('checkcert'), raiseOnError=True
+            checkCert=self._get_option('checkcert'),
+            disableReferral=self._get_option('disablereferral'),
+            raiseOnError=True
         )
         # Reset cache
         self._cached_objects = {}

mylib/oracle.py

@@ -125,7 +125,6 @@ class OracleDB:
             log.debug("Just-try mode : do not really execute SQL query '%s'", sql)
             return True
 
-        cursor = self._conn.cursor()
         try:
             log.debug(
                 'Run SQL query "%s" %s',
@@ -135,10 +134,11 @@ class OracleDB:
                     for key, value in params.items()
                 ]) if params else "without params"
             )
-            if isinstance(params, dict):
-                cursor.execute(sql, **params)
-            else:
-                cursor.execute(sql)
+            with self._conn.cursor() as cursor:
+                if isinstance(params, dict):
+                    cursor.execute(sql, **params)
+                else:
+                    cursor.execute(sql)
             self._conn.commit()
             return True
         except Exception:
@@ -164,7 +164,6 @@ class OracleDB:
         :return: List of selected rows as dict on success, False otherwise
         :rtype: list, bool
         """
-        cursor = self._conn.cursor()
         try:
             log.debug(
                 'Run SQL SELECT query "%s" %s',
@@ -174,12 +173,15 @@ class OracleDB:
                     for key, value in params.items()
                 ]) if params else "without params"
             )
-            if isinstance(params, dict):
-                cursor.execute(sql, **params)
-            else:
-                cursor.execute(sql)
-            cursor.rowfactory = lambda *args: dict(zip([d[0] for d in cursor.description], args))
-            results = cursor.fetchall()
+            with self._conn.cursor() as cursor:
+                if isinstance(params, dict):
+                    cursor.execute(sql, **params)
+                else:
+                    cursor.execute(sql)
+                cursor.rowfactory = lambda *args: dict(
+                    zip([d[0] for d in cursor.description], args)
+                )
+                results = cursor.fetchall()
             return results
         except Exception:
             log.error(

tests/test_oracle.py

@@ -15,8 +15,10 @@ class FakeCXOracleCursor:
         self.expected_return = expected_return
         self.expected_just_try = expected_just_try
         self.expected_exception = expected_exception
+        self.opened = True
 
     def execute(self, sql, **params):
+        assert self.opened
         if self.expected_exception:
             raise Exception("%s.execute(%s, %s): expected exception" % (self, sql, params))
         if self.expected_just_try and not sql.lower().startswith('select '):
@@ -26,8 +28,16 @@ class FakeCXOracleCursor:
         return self.expected_return
 
     def fetchall(self):
+        assert self.opened
         return self.expected_return
 
+    def __enter__(self):
+        self.opened = True
+        return self
+
+    def __exit__(self, *args):
+        self.opened = False
+
     def __repr__(self):
         return "FakeCXOracleCursor(%s, %s, %s, %s)" % (
             self.expected_sql, self.expected_params,