#!/usr/bin/env python3

"""Reformat sqlcmd output to csv

- writes output of the sql query to a temporary file
- saves query as unicode and converts utf-16 to utf-8
- replaces "NULL" with empty cells
"""

import os
import sys
import argparse
import csv
import gzip
from subprocess import run, CalledProcessError
import tempfile


def nonull(row):
    return ['' if x == 'NULL' else x for x in row]


def main(arguments):

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('infile', help="Input file containing an sql command")
    parser.add_argument('-o', '--outfile', default='result.csv',
                        help="Output file name; uses gzip compression if ends with .gz")

    args = parser.parse_args(arguments)
    tempoutfile, tempout = tempfile.mkstemp()
    os.close(tempoutfile)

    with (open(args.infile) as sqlfile,
          tempfile.NamedTemporaryFile('w', delete=False) as sqltemp):
        sqltemp.write('SET NOCOUNT ON;\n\n')
        sqltemp.write(sqlfile.read())

    try:
        # https://learn.microsoft.com/en-us/sql/tools/sqlcmd-utility?view=sql-server-ver16
        cmd = [
            'sqlcmd',
            '-S', 'am-dawg-sql-trt',
            '-i', sqltemp.name,
            '-o', tempout,
            '-s', '|',
            '-k', '2',
            '-E',
            '-W',
            '-m1',
            '-b',
            '-u',
        ]

        run(cmd, check=True)

        if args.outfile.endswith('.gz'):
            opener = gzip.open
        else:
            opener = open

        with (open(tempout, 'r', encoding='utf-16') as tempin,
              opener(args.outfile, 'wt', encoding='utf-8', errors='ignore') as f):
            reader = csv.reader(tempin, delimiter='|')
            headers = next(reader)
            next(reader)  # second row is just dashes
            writer = csv.writer(f, dialect='unix', quoting=csv.QUOTE_MINIMAL)
            writer.writerow(headers)
            writer.writerows((nonull(row) for row in reader))

    except CalledProcessError as err:
        print(err)
        run(['cat', tempout], check=True)
    except Exception as err:
        print(err)
    finally:
        os.remove(tempout)
        os.remove(sqltemp.name)


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))