Fix DB2 LUW reflection, binary binds and XML/DECFLOAT results

ibm_db_sa/base.py

@@ -21,7 +21,7 @@
 """
 import sys
 import sqlalchemy
-import datetime, re
+import datetime, decimal, re
 from sqlalchemy import types as sa_types
 from sqlalchemy import schema as sa_schema
 from sqlalchemy import util
@@ -228,6 +228,30 @@ class DOUBLE(sa_types.Numeric):
     __visit_name__ = 'DOUBLE'
 
 
+class DECFLOAT(sa_types.Numeric):
+    """DB2 DECFLOAT(16) or DECFLOAT(34).
+
+    The ibm_db DBAPI returns DECFLOAT values as str; convert them to Decimal
+    (or float with asdecimal=False) and bind Decimal values as exact text.
+    """
+    __visit_name__ = 'DECFLOAT'
+
+    def __init__(self, precision=34, asdecimal=True):
+        super().__init__(precision=precision, asdecimal=asdecimal)
+
+    def bind_processor(self, dialect):
+        def process(value):
+            return str(value) if isinstance(value, decimal.Decimal) else value
+        return process
+
+    def result_processor(self, dialect, coltype):
+        convert = decimal.Decimal if self.asdecimal else float
+
+        def process(value):
+            return None if value is None else convert(str(value))
+        return process
+
+
 class LONGVARCHAR(sa_types.VARCHAR):
     __visit_name_ = 'LONGVARCHAR'
 
@@ -283,6 +307,9 @@ class XML(sa_types.Text):
     'XML': XML,
     'GRAPHIC': GRAPHIC,
     'VARGRAPHIC': VARGRAPHIC,
+    'DECFLOAT': DECFLOAT,
+    'BINARY': sa_types.BINARY,
+    'VARBINARY': sa_types.VARBINARY,
     'LONGVARGRAPHIC': LONGVARGRAPHIC,
     'DBCLOB': DBCLOB
 }
@@ -301,6 +328,10 @@ def visit_DATE(self, type_, **kw):
        logger.debug(f"Type rendering -> DATE -> {sql}")
        return sql
 
+   @log_entry_exit
+   def visit_DECFLOAT(self, type_, **kw):
+       return "DECFLOAT(%d)" % (type_.precision or 34)
+
    @log_entry_exit
    def visit_TIME(self, type_, **kw):
        sql = "TIME"
@@ -1481,6 +1512,13 @@ def get_indexes(self, connection, table_name, schema=None, **kw):
         logger.debug(f"Indexes fetched -> count={len(indexes)}")
         return indexes
 
+    @log_entry_exit
+    def get_check_constraints(self, connection, table_name, schema=None, **kw):
+        reflect = getattr(self._reflector, "get_check_constraints", None)
+        if reflect is None:
+            raise NotImplementedError()
+        return reflect(connection, table_name, schema=schema, **kw)
+
     @log_entry_exit
     def get_unique_constraints(self, connection, table_name, schema=None, **kw):
         logger.debug(f"Fetching unique constraints -> table={table_name}, schema={schema}")

ibm_db_sa/ibm_db.py

@@ -29,7 +29,7 @@
 from sqlalchemy.engine.url import URL
 from sqlalchemy.exc import ArgumentError
 
-from .base import DB2Dialect, DB2ExecutionContext
+from .base import DB2Dialect, DB2ExecutionContext, DECFLOAT, XML
 from .logger import init_ibmdbsa_logging, log_entry_exit, logger
 
 m = re.match(r"^\s*(\d+)\.(\d+)", SA_VERSION_STR)
@@ -64,6 +64,47 @@ def to_float(value):
            return to_float
 
 
+class _IBM_Binary_ibm_db(sa_types._Binary):
+    """Bind binary values as bytes.
+
+    SQLAlchemy binds binary values through dbapi.Binary, which ibm_db_dbi
+    implements as memoryview. With executemany, ibm_db rejects a memoryview for
+    BINARY, VARBINARY and FOR BIT DATA columns (SQL0302N) and stores its repr
+    text ("<memory at 0x...>") in BLOB columns. bytes work in both paths.
+    """
+
+    def bind_processor(self, dialect):
+        def process(value):
+            return None if value is None else bytes(value)
+        return process
+
+
+# DB2 stores XML parsed, without a declaration. On fetch the CLI serializes it
+# with a byte order mark and a UTF-16 declaration, even for a document stored
+# with a UTF-8 declaration, which is wrong for a Python str.
+_CLI_XML_PREFIX = re.compile(r'\A\ufeff?(?:<\?xml version="1\.0" encoding="UTF-16" \?>)?')
+
+
+class _IBM_XML_ibm_db(XML):
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            if isinstance(value, str):
+                return _CLI_XML_PREFIX.sub("", value, count=1)
+            return value
+        return process
+
+
+_DISCONNECT_MESSAGES = (
+    'Connection is not active',
+    'connection is no longer active',
+    'Connection Resource cannot be found',
+    'SQL30081N',
+    'CLI0108E',
+    'CLI0106E',
+    'SQL1224N',
+)
+
+
 class DB2ExecutionContext_ibm_db(DB2ExecutionContext):
     _callproc_result = None
     _out_parameters = None
@@ -128,7 +169,11 @@ class DB2Dialect_ibm_db(DB2Dialect):
     colspecs = util.update_copy(
         DB2Dialect.colspecs,
         {
-            sa_types.Numeric: _IBM_Numeric_ibm_db
+            sa_types.Numeric: _IBM_Numeric_ibm_db,
+            sa_types._Binary: _IBM_Binary_ibm_db,
+            # DECFLOAT is a Numeric but keeps its own processors.
+            DECFLOAT: DECFLOAT,
+            XML: _IBM_XML_ibm_db,
         }
     )
 
@@ -303,27 +348,25 @@ def _get_default_schema_name(self, connection):
         logger.debug("Normalized schema: %s", normalized_schema_name)
         return normalized_schema_name
 
-    # Checks if the DB_API driver error indicates an invalid connection
+    # Checks if the DB_API driver error indicates an invalid connection. A
+    # connection lost while fetching surfaces as the base ibm_db_dbi.Error.
     @log_entry_exit
     def is_disconnect(self, ex, connection, cursor):
-        logger.debug("Checking if exception indicates disconnect")
-        logger.debug("Exception received: %s", ex)
-        if isinstance(ex, (self.dbapi.ProgrammingError,
-                           self.dbapi.OperationalError)):
-            connection_errors = ('Connection is not active',
-                                 'connection is no longer active',
-                                 'Connection Resource cannot be found',
-                                 'SQL30081N',
-                                 'CLI0108E',
-                                 'CLI0106E',
-                                 'SQL1224N')
-            for err_msg in connection_errors:
+        if isinstance(ex, self.dbapi.Error):
+            for err_msg in _DISCONNECT_MESSAGES:
                 if err_msg in str(ex):
                     logger.debug("Disconnect detected due to error: %s", err_msg)
                     return True
-        else:
-            logger.debug("Exception type does not indicate disconnect")
         return False
 
+    # After a server restart or network failure ibm_db_dbi raises CLI0106E
+    # ("Connection is closed") from close(); the pool then logged an error for
+    # every connection it discarded.
+    def do_close(self, dbapi_connection):
+        try:
+            dbapi_connection.close()
+        except self.dbapi.Error as err:
+            if "CLI0106E" not in str(err):
+                raise
 
 dialect = DB2Dialect_ibm_db