From 91573cb807f6399cebb8ce8d19fad1758578811c Mon Sep 17 00:00:00 2001
From: Matt Gilman <matt.c.gilman@gmail.com>
Date: Mon, 23 Nov 2015 08:45:24 -0500
Subject: [PATCH] NIFI-655: - Fixing issue detecting the presence of DN column

---
 .../nifi/admin/AuditDataSourceFactoryBean.java | 18 ++++--------------
 .../nifi/admin/UserDataSourceFactoryBean.java  | 18 ++++--------------
 2 files changed, 8 insertions(+), 28 deletions(-)

diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-administration/src/main/java/org/apache/nifi/admin/AuditDataSourceFactoryBean.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-administration/src/main/java/org/apache/nifi/admin/AuditDataSourceFactoryBean.java
index 7e6879916c..87cd420fd3 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-administration/src/main/java/org/apache/nifi/admin/AuditDataSourceFactoryBean.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-administration/src/main/java/org/apache/nifi/admin/AuditDataSourceFactoryBean.java
@@ -19,7 +19,6 @@ package org.apache.nifi.admin;
 import java.io.File;
 import java.sql.Connection;
 import java.sql.ResultSet;
-import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
 import org.apache.commons.lang3.StringUtils;
@@ -176,10 +175,11 @@ public class AuditDataSourceFactoryBean implements FactoryBean {
                     statement.execute(CREATE_PURGE_DETAILS_TABLE);
                 } else {
                     logger.info("Existing database found and connected to at: " + databaseUrl);
+                    RepositoryUtils.closeQuietly(rs);
 
-                    // get the RS metadata to see if we need to transform the table
-                    final ResultSetMetaData rsMetadata = rs.getMetaData();
-                    if (hasDnColumn(rsMetadata)) {
+                    // check if the DN column exists to see if we need to transform the table
+                    rs = connection.getMetaData().getColumns(null, null, "ACTION", "USER_DN");
+                    if (rs.next()) {
                         statement.execute(RENAME_DN_COLUMN);
                         statement.execute(RESIZE_IDENTITY_COLUMN);
                         statement.execute(RESIZE_USER_NAME_COLUMN);
@@ -201,16 +201,6 @@ public class AuditDataSourceFactoryBean implements FactoryBean {
         return connectionPool;
     }
 
-    private boolean hasDnColumn(final ResultSetMetaData rsMetadata) throws SQLException {
-        boolean hasDn = false;
-        for (int i = 1; i <= rsMetadata.getColumnCount() && !hasDn; i++) {
-            if ("USER_DN".equals(rsMetadata.getColumnName(i))) {
-                hasDn = true;
-            }
-        }
-        return hasDn;
-    }
-
     @Override
     public Class getObjectType() {
         return JdbcConnectionPool.class;
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-administration/src/main/java/org/apache/nifi/admin/UserDataSourceFactoryBean.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-administration/src/main/java/org/apache/nifi/admin/UserDataSourceFactoryBean.java
index acd34495ac..6d8566e8ad 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-administration/src/main/java/org/apache/nifi/admin/UserDataSourceFactoryBean.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-administration/src/main/java/org/apache/nifi/admin/UserDataSourceFactoryBean.java
@@ -19,7 +19,6 @@ package org.apache.nifi.admin;
 import java.io.File;
 import java.sql.Connection;
 import java.sql.ResultSet;
-import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.HashSet;
@@ -151,10 +150,11 @@ public class UserDataSourceFactoryBean implements FactoryBean {
                     statement.execute(INSERT_ANONYMOUS_USER);
                 } else {
                     logger.info("Existing database found and connected to at: " + databaseUrl);
+                    RepositoryUtils.closeQuietly(rs);
 
-                    // get the RS metadata to see if we need to transform the table
-                    final ResultSetMetaData rsMetadata = rs.getMetaData();
-                    if (hasDnColumn(rsMetadata)) {
+                    // if the DN column exists, transform the table
+                    rs = connection.getMetaData().getColumns(null, null, "USER", "DN");
+                    if (rs.next()) {
                         statement.execute(RENAME_DN_COLUMN);
                         statement.execute(RESIZE_IDENTITY_COLUMN);
                         statement.execute(RESIZE_USER_NAME_COLUMN);
@@ -184,16 +184,6 @@ public class UserDataSourceFactoryBean implements FactoryBean {
         return connectionPool;
     }
 
-    private boolean hasDnColumn(final ResultSetMetaData rsMetadata) throws SQLException {
-        boolean hasDn = false;
-        for (int i = 1; i <= rsMetadata.getColumnCount() && !hasDn; i++) {
-            if ("DN".equals(rsMetadata.getColumnName(i))) {
-                hasDn = true;
-            }
-        }
-        return hasDn;
-    }
-
     private String getDatabaseUrl(File databaseFile) {
         String databaseUrl = "jdbc:h2:" + databaseFile + ";AUTOCOMMIT=OFF;DB_CLOSE_ON_EXIT=FALSE;LOCK_MODE=3";
         String databaseUrlAppend = properties.getProperty(NiFiProperties.H2_URL_APPEND);