package com.hortonworks.spark.sql.hive.llap;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.execution.datasources.OutputWriter;
import org.apache.spark.sql.execution.datasources.orc.OrcOutputWriter;
import org.apache.spark.sql.sources.v2.writer.DataWriter;
import org.apache.spark.sql.sources.v2.writer.DataWriterFactory;
import org.apache.spark.sql.sources.v2.writer.WriterCommitMessage;
import org.apache.spark.sql.types.StructType;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static com.hortonworks.spark.sql.hive.llap.MockHiveWarehouseConnector.testVector;
import static org.junit.Assert.assertEquals;

public class MockWriteSupport {

  public static class MockHiveWarehouseDataSourceWriter extends HiveWarehouseDataSourceWriter {

    public MockHiveWarehouseDataSourceWriter(Map<String, String> options, String jobId, StructType schema, Path path,
        Configuration conf) {
      super(options, jobId, schema, path, conf);
    }

    @Override
    public DataWriterFactory<InternalRow> createInternalRowWriterFactory() {
      ByteArrayOutputStream confByteArrayStream = new ByteArrayOutputStream();
      byte[] confBytes;
      try(DataOutputStream confByteData = new DataOutputStream(confByteArrayStream)) {
        conf.write(confByteData);
        confBytes = confByteArrayStream.toByteArray();
      } catch(Exception e) {
        throw new RuntimeException(e);
      }
      return new MockHiveWarehouseDataWriterFactory(jobId, schema, path.toString(), confBytes);
    }

    @Override public void commit(WriterCommitMessage[] messages) {

    }
  }

  public static class MockHiveWarehouseDataWriterFactory extends HiveWarehouseDataWriterFactory {

    public MockHiveWarehouseDataWriterFactory(String jobId, StructType schema, String path,
        byte[] confBytes) {
      super(jobId, schema, path, confBytes);
    }

    protected DataWriter<InternalRow> getDataWriter(Configuration conf, String jobId,
        StructType schema, int partitionId, int attemptNumber,
        FileSystem fs, Path filePath) {
      return new MockHiveWarehouseDataWriter(conf, jobId, schema, partitionId, attemptNumber, fs, filePath);
    }

  }

  public static class MockHiveWarehouseDataWriter extends HiveWarehouseDataWriter {

    public MockHiveWarehouseDataWriter(Configuration conf, String jobId, StructType schema, int partitionId,
        int attemptNumber, FileSystem fs, Path filePath) {
      super(conf, jobId, schema, partitionId, attemptNumber, fs, filePath);
    }

    @Override
    protected OutputWriter getOutputWriter(String path, StructType schema, TaskAttemptContext tac) {
      return new MockOutputWriter(path, schema, tac);
    }

  }

  public static class MockOutputWriter extends OutputWriter {

    public static List<InternalRow> rowBuffer = new ArrayList<>();
    public static boolean closed = false;

    public MockOutputWriter(String path, StructType schema, TaskAttemptContext tac) {
    }

    @Override public void write(InternalRow row) {
      rowBuffer.add(row);
    }

    @Override public void close() {
      MockHiveWarehouseConnector.writeOutputBuffer.put("TestWriteSupport", rowBuffer);
    }
  }

}

@@ -35,24 +35,105 @@
 hive.llap;%0A%0A
+import com.hortonworks.spark.sql.hive.llap.util.SerializableHadoopConfiguration;%0A
 import org.a
@@ -1425,363 +1425,8 @@
 ) %7B%0A
-      ByteArrayOutputStream confByteArrayStream = new ByteArrayOutputStream();%0A      byte%5B%5D confBytes;%0A      try(DataOutputStream confByteData = new DataOutputStream(confByteArrayStream)) %7B%0A        conf.write(confByteData);%0A        confBytes = confByteArrayStream.toByteArray();%0A      %7D catch(Exception e) %7B%0A        throw new RuntimeException(e);%0A      %7D%0A
     
@@ -1496,30 +1496,51 @@
 path
-.toString(), confBytes
+, new SerializableHadoopConfiguration(conf)
 );%0A 
@@ -1806,45 +1806,55 @@
 ma, 
-String path,%0A        byte%5B%5D confBytes
+Path path, SerializableHadoopConfiguration conf
 ) %7B%0A
@@ -1894,13 +1894,8 @@
 conf
-Bytes
 );%0A 

Fix tests failures caused by dependency on inaccessible Spark class SerializableConfiguration