PostgreSQLOidDataType.java

  1. package org.dbunit.ext.postgresql;

  2. import org.dbunit.dataset.datatype.BytesDataType;
  3. import org.dbunit.dataset.datatype.TypeCastException;
  4. import org.postgresql.PGConnection;
  5. import org.postgresql.largeobject.LargeObject;
  6. import org.postgresql.largeobject.LargeObjectManager;
  7. import org.slf4j.Logger;
  8. import org.slf4j.LoggerFactory;

  9. import java.io.ByteArrayInputStream;
  10. import java.sql.*;

  11. public class PostgreSQLOidDataType
  12.         extends BytesDataType {

  13.     /**
  14.      * Logger for this class
  15.      */
  16.     private static final Logger logger = LoggerFactory.getLogger(PostgreSQLOidDataType.class);

  17.     public PostgreSQLOidDataType() {
  18.         super("OID", Types.BIGINT);
  19.     }

  20.     @Override
  21.     public Object getSqlValue(final int column, final ResultSet resultSet)
  22.             throws SQLException,
  23.             TypeCastException
  24.     {
  25.         logger.debug("getSqlValue(column={}, resultSet={}) - start", column, resultSet);

  26.         Statement statement = resultSet.getStatement();
  27.         Connection connection = statement.getConnection();
  28.         boolean autoCommit = connection.getAutoCommit();
  29.         // kinda ugly
  30.         connection.setAutoCommit(false);

  31.         try {
  32.             PGConnection pgConnection = connection.unwrap(PGConnection.class);
  33.             LargeObjectManager lobj = pgConnection.getLargeObjectAPI();

  34.             long oid = resultSet.getLong(column);
  35.             if (oid == 0) {
  36.                 logger.debug("'oid' is zero, the data is NULL.");
  37.                 return null;
  38.             }
  39.             LargeObject obj = lobj.open(oid, LargeObjectManager.READ);
  40.             // If lobj.open() throws an exception, it means something wrong with the OID / table.
  41.             // So to be accurate, we don't catch this exception but let it propagate.
  42.             // Swallowing the exception silently indeed hides the problem, which is the wrong behavior
  43. //            try {
  44. //                obj = lobj.open(oid, LargeObjectManager.READ);
  45. //            } catch (SQLException ex) {
  46. //                logger.error("Failed to open oid {} for Large Object reading.", oid);
  47. //                logger.error("Exception: {}", ex.getMessage());
  48. //                logger.error("Returning null instead of bailing out");
  49. //                return null;
  50. //            }

  51.             // Read the data
  52.             byte buf[] = new byte[obj.size()];
  53.             obj.read(buf, 0, obj.size());
  54.             // Close the object
  55.             obj.close();

  56.             return buf;
  57.         } finally {
  58.             connection.setAutoCommit(autoCommit);
  59.         }
  60.     }

  61.     @Override
  62.     public void setSqlValue(final Object value, final int column, final PreparedStatement statement)
  63.             throws SQLException, TypeCastException
  64.     {
  65.         logger.debug("setSqlValue(value={}, column={}, statement={}) - start",
  66.             value, column, statement);

  67.         Connection connection = statement.getConnection();
  68.         boolean autoCommit = connection.getAutoCommit();
  69.         // kinda ugly
  70.         connection.setAutoCommit(false);

  71.         try {
  72.             // Get the Large Object Manager to perform operations with
  73.             LargeObjectManager lobj = (connection.unwrap(PGConnection.class)).getLargeObjectAPI();

  74.             // Create a new large object
  75.             long oid = lobj.createLO(LargeObjectManager.READ | LargeObjectManager.WRITE);

  76.             // Open the large object for writing
  77.             LargeObject obj = lobj.open(oid, LargeObjectManager.WRITE);

  78.             // Now open the file
  79.             ByteArrayInputStream bis = new ByteArrayInputStream((byte[]) super.typeCast(value));

  80.             // Copy the data from the file to the large object
  81.             byte buf[] = new byte[2048];
  82.             int s = 0;
  83.             while ((s = bis.read(buf, 0, 2048)) > 0) {
  84.                 obj.write(buf, 0, s);
  85.             }

  86.             // Close the large object
  87.             obj.close();

  88.             statement.setLong(column, oid);
  89.         } finally {
  90.             connection.setAutoCommit(autoCommit);
  91.         }
  92.     }
  93. }