Close read connection after write (#314)

connector/src/main/scala/com/vertica/spark/datasource/core/DSWriter.scala

@@ -94,6 +94,9 @@ class DSWriter(config: WriteConfig, uniqueId: String, pipeFactory: VerticaPipeFa
   }
 
   def commitRows(): ConnectorResult[Unit] = {
-    pipe.commit()
+    val ret = pipe.commit()
+    // Ensure all connections are closed, including read connections used by the write operation
+    val _ = pipeFactory.closeJdbcLayers()
+    ret
   }
 }

connector/src/main/scala/com/vertica/spark/datasource/core/factory/VerticaPipeFactory.scala

@@ -30,6 +30,8 @@ trait VerticaPipeFactoryInterface {
   def getReadPipe(config: ReadConfig): VerticaPipeInterface with VerticaPipeReadInterface
 
   def getWritePipe(config: WriteConfig): VerticaPipeInterface with VerticaPipeWriteInterface
+
+  def closeJdbcLayers()
 }
 
 /**
@@ -82,4 +84,15 @@ object VerticaPipeFactory extends VerticaPipeFactoryInterface {
     }
   }
 
+  override def closeJdbcLayers(): Unit = {
+    readLayer match {
+      case Some(layer) => val _ = layer.close
+      case None =>
+    }
+    writeLayer match {
+      case Some(layer) => val _ = layer.close
+      case None =>
+    }
+  }
+
 }

connector/src/test/scala/com/vertica/spark/datasource/core/DSWriterTest.scala

@@ -151,6 +151,7 @@ class DSWriterTest extends AnyFlatSpec with BeforeAndAfterAll with MockFactory {
     (pipe.commit _).expects().returning(Right())
     val pipeFactory = mock[VerticaPipeFactoryInterface]
     (pipeFactory.getWritePipe _).expects(*).returning(pipe)
+    (pipeFactory.closeJdbcLayers _).expects().returning(())
 
     val writer = new DSWriter(config, "unique-id", pipeFactory)
     checkResult(writer.commitRows())