//
// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
//
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using Microsoft.Kusto.ServiceLayer.QueryExecution.Contracts;
namespace Microsoft.Kusto.ServiceLayer.QueryExecution.DataStorage
{
///
/// Writer for writing rows of results to a CSV file
///
public class SaveAsCsvFileStreamWriter : SaveAsStreamWriter
{
#region Member Variables
private readonly SaveResultsAsCsvRequestParams saveParams;
private bool headerWritten;
#endregion
///
/// Constructor, stores the CSV specific request params locally, chains into the base
/// constructor
///
/// FileStream to access the CSV file output
/// CSV save as request parameters
public SaveAsCsvFileStreamWriter(Stream stream, SaveResultsAsCsvRequestParams requestParams)
: base(stream, requestParams)
{
saveParams = requestParams;
}
///
/// Writes a row of data as a CSV row. If this is the first row and the user has requested
/// it, the headers for the column will be emitted as well.
///
/// The data of the row to output to the file
///
/// The entire list of columns for the result set. They will be filtered down as per the
/// request params.
///
public override void WriteRow(IList row, IList columns)
{
char delimiter = ',';
if(!string.IsNullOrEmpty(saveParams.Delimiter))
{
// first char in string
delimiter = saveParams.Delimiter[0];
}
string lineSeperator = Environment.NewLine;
if(!string.IsNullOrEmpty(saveParams.LineSeperator))
{
lineSeperator = saveParams.LineSeperator;
}
char textIdentifier = '"';
if(!string.IsNullOrEmpty(saveParams.TextIdentifier))
{
// first char in string
textIdentifier = saveParams.TextIdentifier[0];
}
Encoding.RegisterProvider(CodePagesEncodingProvider.Instance);
int codepage;
Encoding encoding;
try
{
if(int.TryParse(saveParams.Encoding, out codepage))
{
encoding = Encoding.GetEncoding(codepage);
}
else
{
encoding = Encoding.GetEncoding(saveParams.Encoding);
}
}
catch
{
// Fallback encoding when specified codepage is invalid
encoding = Encoding.GetEncoding("utf-8");
}
// Write out the header if we haven't already and the user chose to have it
if (saveParams.IncludeHeaders && !headerWritten)
{
// Build the string
var selectedColumns = columns.Skip(ColumnStartIndex ?? 0).Take(ColumnCount ?? columns.Count)
.Select(c => EncodeCsvField(c.ColumnName, delimiter, textIdentifier) ?? string.Empty);
string headerLine = string.Join(delimiter, selectedColumns);
// Encode it and write it out
byte[] headerBytes = encoding.GetBytes(headerLine + lineSeperator);
FileStream.Write(headerBytes, 0, headerBytes.Length);
headerWritten = true;
}
// Build the string for the row
var selectedCells = row.Skip(ColumnStartIndex ?? 0)
.Take(ColumnCount ?? columns.Count)
.Select(c => EncodeCsvField(c.DisplayValue, delimiter, textIdentifier));
string rowLine = string.Join(delimiter, selectedCells);
// Encode it and write it out
byte[] rowBytes = encoding.GetBytes(rowLine + lineSeperator);
FileStream.Write(rowBytes, 0, rowBytes.Length);
}
///
/// Encodes a single field for inserting into a CSV record. The following rules are applied:
///
/// - All double quotes (") are replaced with a pair of consecutive double quotes
///
/// The entire field is also surrounded by a pair of double quotes if any of the following conditions are met:
///
/// - The field begins or ends with a space
/// - The field begins or ends with a tab
/// - The field contains the ListSeparator string
/// - The field contains the '\n' character
/// - The field contains the '\r' character
/// - The field contains the '"' character
///
///
/// The field to encode
/// The CSV encoded version of the original field
internal static string EncodeCsvField(string field, char delimiter, char textIdentifier)
{
string strTextIdentifier = textIdentifier.ToString();
// Special case for nulls
if (field == null)
{
return "NULL";
}
// Whether this field has special characters which require it to be embedded in quotes
bool embedInQuotes = field.IndexOfAny(new[] { delimiter, '\r', '\n', textIdentifier }) >= 0 // Contains special characters
|| field.StartsWith(" ") || field.EndsWith(" ") // Start/Ends with space
|| field.StartsWith("\t") || field.EndsWith("\t"); // Starts/Ends with tab
//Replace all quotes in the original field with double quotes
string ret = field.Replace(strTextIdentifier, strTextIdentifier + strTextIdentifier);
if (embedInQuotes)
{
ret = strTextIdentifier + $"{ret}" + strTextIdentifier;
}
return ret;
}
}
}