import json
import csv
import sys
import os
import re
import codecs
import logging
from logging.config import dictConfig

import click
import yaml
from sqlalchemy import create_engine
from jsontableschema_sql import Storage
from smart_open import smart_open

from . import postgres
from . import carto

csv.field_size_limit(sys.maxsize)

def get_logger(logging_config):
    try:
        with open(logging_config) as file:
            config = yaml.load(file)
        dictConfig(config)
    except:
        FORMAT = '[%(asctime)-15s] %(levelname)s [%(name)s] %(message)s'
        logging.basicConfig(format=FORMAT, level=logging.INFO, stream=sys.stderr)

    logger = logging.getLogger('the_el')

    def exception_handler(type, value, tb):
        logger.exception("Uncaught exception: {}".format(str(value)), exc_info=(type, value, tb))

    sys.excepthook = exception_handler

    return logger

@click.group()
def main():
    pass

def get_connection_string(connection_string):
    connection_string = os.getenv('CONNECTION_STRING', connection_string)
    if connection_string == None:
        raise Exception('`CONNECTION_STRING` environment variable or `--connection-string` option required')
    return connection_string

def create_storage_adaptor(connection_string, db_schema, geometry_support, from_srid=None, to_srid=None):
    engine = create_engine(connection_string)
    storage = Storage(engine, dbschema=db_schema, geometry_support=geometry_support, from_srid=from_srid, to_srid=to_srid, views=True)
    return engine, storage

def fopen(file, mode='r'):
    if file == None:
        if mode == 'r':
            return sys.stdin
        elif mode == 'w':
            return sys.stdout
    else:
        return smart_open(file, mode=mode)

def get_table_schema(table_schema_path):
    with fopen(table_schema_path) as file:
        contents = file.read()
        if not isinstance(contents, str):
             contents = contents.decode('utf-8')
        return json.loads(contents)

@main.command()
@click.argument('table_name')
@click.option('--connection-string')
@click.option('-o','--output-file')
@click.option('--db-schema')
@click.option('--geometry-support')
def describe_table(table_name, connection_string, output_file, db_schema, geometry_support):
    connection_string = get_connection_string(connection_string)

    engine, storage = create_storage_adaptor(connection_string, db_schema, geometry_support)
    descriptor = storage.describe(table_name)

    with fopen(output_file, mode='w') as file:
        json.dump(descriptor, file)

@main.command()
@click.argument('table_name')
@click.argument('table_schema_path')
@click.option('--connection-string')
@click.option('--db-schema')
@click.option('--indexes-fields')
@click.option('--geometry-support')
@click.option('--if-not-exists', is_flag=True, default=False)
@click.option('--logging-config', default='logging_config.conf')
def create_table(table_name,
                 table_schema_path,
                 connection_string,
                 db_schema,
                 indexes_fields,
                 geometry_support,
                 if_not_exists,
                 logging_config):
    logger = get_logger(logging_config)

    table_schema = get_table_schema(table_schema_path)

    if indexes_fields != None:
        indexes_fields = indexes_fields.split(',')

    if re.match(carto.carto_connection_string_regex, connection_string) != None:
        load_postgis = geometry_support == 'postgis'
        logger.info('{} - Creating table using Carto'.format(table_name))
        return carto.create_table(logger, table_name, load_postgis, table_schema, if_not_exists, indexes_fields, connection_string)

    connection_string = get_connection_string(connection_string)

    engine, storage = create_storage_adaptor(connection_string, db_schema, geometry_support)

    logger.info('{} - Creating table using SQLAlchemy'.format(table_name))
    storage.create(table_name, table_schema, indexes_fields=indexes_fields)

@main.command()
@click.argument('table_name')
@click.option('--table-schema-path')
@click.option('--connection-string')
@click.option('-f','--input-file')
@click.option('--db-schema')
@click.option('--geometry-support')
@click.option('--from-srid')
@click.option('--skip-headers', is_flag=True)
@click.option('--indexes-fields')
@click.option('--upsert', is_flag=True)
@click.option('--truncate/--no-truncate', is_flag=True, default=False)
@click.option('--logging-config', default='logging_config.conf')
def write(table_name,
          table_schema_path,
          connection_string,
          input_file,
          db_schema,
          geometry_support,
          from_srid,
          skip_headers,
          indexes_fields,
          upsert,
          truncate,
          logging_config):
    logger = get_logger(logging_config)

    table_schema = get_table_schema(table_schema_path)

    ## TODO: csv settings? use Frictionless Data csv standard?
    ## TODO: support line delimted json?
    with fopen(input_file) as file:
        rows = csv.reader(file)

        if skip_headers:
            next(rows)

        if re.match(carto.carto_connection_string_regex, connection_string) != None:
            load_postgis = geometry_support == 'postgis'

            if indexes_fields != None:
                indexes_fields = indexes_fields.split(',')

            logger.info('{} - Writing to table using Carto'.format(table_name))

            carto.load(logger,
                       db_schema,
                       table_name,
                       load_postgis,
                       table_schema,
                       connection_string,
                       rows,
                       indexes_fields,
                       truncate)
        else:
            connection_string = get_connection_string(connection_string)

            engine, storage = create_storage_adaptor(connection_string, db_schema, geometry_support, from_srid=from_srid)

            ## TODO: truncate? carto does. Makes this idempotent

            logger.info('{} - Writing to table using SQLAlchemy'.format(table_name))

            if table_schema_path != None:
                table_schema = get_table_schema(table_schema_path)
                storage.describe(table_name, descriptor=table_schema)
            else:
                storage.describe(table_name)

            if upsert:
                postgres.upsert(engine, db_schema, table_name, table_schema, rows)
            elif geometry_support == None and engine.dialect.driver == 'psycopg2':
                postgres.copy_from(engine, table_name, table_schema, rows)
            else:
                storage.write(table_name, rows)

@main.command()
@click.argument('table_name')
@click.option('--connection-string')
@click.option('-o','--output-file')
@click.option('--db-schema')
@click.option('--geometry-support')
@click.option('--from-srid')
@click.option('--to-srid')
@click.option('--logging-config', default='logging_config.conf')
def read(table_name, connection_string, output_file, db_schema, geometry_support, from_srid, to_srid, logging_config):
    logger = get_logger(logging_config)

    connection_string = get_connection_string(connection_string)

    engine, storage = create_storage_adaptor(connection_string, db_schema, geometry_support, from_srid=from_srid, to_srid=to_srid)

    ## TODO: csv settings? use Frictionless Data csv standard?
    ## TODO: support line delimited json?
    with fopen(output_file, mode='w') as file:
        writer = csv.writer(file)

        descriptor = storage.describe(table_name)
        fields = map(lambda x: x['name'], descriptor['fields'])
        writer.writerow(fields)

        if geometry_support == None and engine.dialect.driver == 'psycopg2':
            postgres.copy_to(engine, table_name, file)
        else:
            for row in storage.iter(table_name):
                row_out = []
                for field in row:
                    if isinstance(field, dict) or isinstance(field, list):
                        field = json.dumps(field)
                    row_out.append(field)
                writer.writerow(row_out)

@main.command()
@click.argument('new_table_name')
@click.argument('old_table_name')
@click.option('--connection-string')
@click.option('--db-schema')
@click.option('--select-users', help='Users to grant SELECT on updated table')
@click.option('--logging-config', default='logging_config.conf')
def swap_table(new_table_name, old_table_name, connection_string, db_schema, select_users, logging_config):
    logger = get_logger(logging_config)

    if re.match(carto.carto_connection_string_regex, connection_string) != None:
        if select_users != None:
            select_users = select_users.split(',')
        else:
            select_users = []
        logger.info('Swapping tables using Carto: {} - {}'.format(new_table_name, old_table_name))
        return carto.swap_table(logger, db_schema, new_table_name, old_table_name, select_users, connection_string)

    connection_string = get_connection_string(connection_string)
    engine = create_engine(connection_string)
 
    if engine.dialect.driver == 'psycopg2':
        logger.info('Swapping tables using psycopg2: {} - {}'.format(new_table_name, old_table_name))
        conn = engine.raw_connection()
        try:
            with conn.cursor() as cur:
                sql = 'ALTER TABLE "{}" RENAME TO "{}_old";'.format(old_table_name, old_table_name) +\
                      'ALTER TABLE "{}" RENAME TO "{}";'.format(new_table_name, old_table_name) +\
                      'DROP TABLE "{}_old";'.format(old_table_name)
                cur.execute(sql)
            conn.commit()
        except:
            conn.rollback()
            raise
        conn.close()
    elif engine.dialect.driver == 'cx_oracle':
        logger.info('Swapping tables using cx_Oracle: {} - {}'.format(new_table_name, old_table_name))
        conn = engine.connect()
        if select_users != None:
            select_users = select_users.split(',')
        else:
            select_users = []
        grants_sql = []
        for user in select_users:
            grants_sql.append('GRANT SELECT ON {} TO {}'.format(old_table_name, user.strip()))

        # Oracle does not allow table modification within a transaction, so make individual transactions:
        sql1 = 'ALTER TABLE {} RENAME TO {}_old'.format(old_table_name, old_table_name)
        sql2 = 'ALTER TABLE {} RENAME TO {}'.format(new_table_name, old_table_name)
        sql3 = 'DROP TABLE {}_old'.format(old_table_name)

        try:
            conn.execute(sql1)
        except:
            logger.error("Could not rename {} table. Does it exist?".format(old_table_name))
            raise
        try:
            conn.execute(sql2)
        except:
            logger.error("Could not rename {} table. Does it exist?".format(new_table_name))
            rb_sql = 'ALTER TABLE {}_old RENAME TO {}'.format(old_table_name, old_table_name)
            conn.execute(rb_sql)
            raise
        try:
            conn.execute(sql3)
        except:
            logger.error("Could not drop {}_old table. Do you have permission?".format(old_table_name))
            rb_sql1 = 'DROP TABLE {}'.format(old_table_name)
            conn.execute(rb_sql1)
            rb_sql2 = 'ALTER TABLE {}_old RENAME TO {}'.format(old_table_name, old_table_name)
            conn.execute(rb_sql2)
            raise
        try:
            for sql in grants_sql:
                conn.execute(sql)
        except:
            logger.error("Could not grant all permissions to {}.".format(old_table_name))
            raise
    else:
        raise Exception('`{}` not supported by swap_table'.format(engine.dialect.driver))

[((303, 336), 'csv.field_size_limit', 'csv.field_size_limit', (['sys.maxsize'], {}), '(sys.maxsize)\n', (323, 336), False, 'import csv\n'), ((899, 912), 'click.group', 'click.group', ([], {}), '()\n', (910, 912), False, 'import click\n'), ((2014, 2042), 'click.argument', 'click.argument', (['"""table_name"""'], {}), "('table_name')\n", (2028, 2042), False, 'import click\n'), ((2044, 2079), 'click.option', 'click.option', (['"""--connection-string"""'], {}), "('--connection-string')\n", (2056, 2079), False, 'import click\n'), ((2081, 2116), 'click.option', 'click.option', (['"""-o"""', '"""--output-file"""'], {}), "('-o', '--output-file')\n", (2093, 2116), False, 'import click\n'), ((2117, 2144), 'click.option', 'click.option', (['"""--db-schema"""'], {}), "('--db-schema')\n", (2129, 2144), False, 'import click\n'), ((2146, 2180), 'click.option', 'click.option', (['"""--geometry-support"""'], {}), "('--geometry-support')\n", (2158, 2180), False, 'import click\n'), ((2581, 2609), 'click.argument', 'click.argument', (['"""table_name"""'], {}), "('table_name')\n", (2595, 2609), False, 'import click\n'), ((2611, 2646), 'click.argument', 'click.argument', (['"""table_schema_path"""'], {}), "('table_schema_path')\n", (2625, 2646), False, 'import click\n'), ((2648, 2683), 'click.option', 'click.option', (['"""--connection-string"""'], {}), "('--connection-string')\n", (2660, 2683), False, 'import click\n'), ((2685, 2712), 'click.option', 'click.option', (['"""--db-schema"""'], {}), "('--db-schema')\n", (2697, 2712), False, 'import click\n'), ((2714, 2746), 'click.option', 'click.option', (['"""--indexes-fields"""'], {}), "('--indexes-fields')\n", (2726, 2746), False, 'import click\n'), ((2748, 2782), 'click.option', 'click.option', (['"""--geometry-support"""'], {}), "('--geometry-support')\n", (2760, 2782), False, 'import click\n'), ((2784, 2844), 'click.option', 'click.option', (['"""--if-not-exists"""'], {'is_flag': '(True)', 'default': '(False)'}), "('--if-not-exists', is_flag=True, default=False)\n", (2796, 2844), False, 'import click\n'), ((2846, 2909), 'click.option', 'click.option', (['"""--logging-config"""'], {'default': '"""logging_config.conf"""'}), "('--logging-config', default='logging_config.conf')\n", (2858, 2909), False, 'import click\n'), ((4023, 4051), 'click.argument', 'click.argument', (['"""table_name"""'], {}), "('table_name')\n", (4037, 4051), False, 'import click\n'), ((4053, 4088), 'click.option', 'click.option', (['"""--table-schema-path"""'], {}), "('--table-schema-path')\n", (4065, 4088), False, 'import click\n'), ((4090, 4125), 'click.option', 'click.option', (['"""--connection-string"""'], {}), "('--connection-string')\n", (4102, 4125), False, 'import click\n'), ((4127, 4161), 'click.option', 'click.option', (['"""-f"""', '"""--input-file"""'], {}), "('-f', '--input-file')\n", (4139, 4161), False, 'import click\n'), ((4162, 4189), 'click.option', 'click.option', (['"""--db-schema"""'], {}), "('--db-schema')\n", (4174, 4189), False, 'import click\n'), ((4191, 4225), 'click.option', 'click.option', (['"""--geometry-support"""'], {}), "('--geometry-support')\n", (4203, 4225), False, 'import click\n'), ((4227, 4254), 'click.option', 'click.option', (['"""--from-srid"""'], {}), "('--from-srid')\n", (4239, 4254), False, 'import click\n'), ((4256, 4300), 'click.option', 'click.option', (['"""--skip-headers"""'], {'is_flag': '(True)'}), "('--skip-headers', is_flag=True)\n", (4268, 4300), False, 'import click\n'), ((4302, 4334), 'click.option', 'click.option', (['"""--indexes-fields"""'], {}), "('--indexes-fields')\n", (4314, 4334), False, 'import click\n'), ((4336, 4374), 'click.option', 'click.option', (['"""--upsert"""'], {'is_flag': '(True)'}), "('--upsert', is_flag=True)\n", (4348, 4374), False, 'import click\n'), ((4376, 4445), 'click.option', 'click.option', (['"""--truncate/--no-truncate"""'], {'is_flag': '(True)', 'default': '(False)'}), "('--truncate/--no-truncate', is_flag=True, default=False)\n", (4388, 4445), False, 'import click\n'), ((4447, 4510), 'click.option', 'click.option', (['"""--logging-config"""'], {'default': '"""logging_config.conf"""'}), "('--logging-config', default='logging_config.conf')\n", (4459, 4510), False, 'import click\n'), ((6711, 6739), 'click.argument', 'click.argument', (['"""table_name"""'], {}), "('table_name')\n", (6725, 6739), False, 'import click\n'), ((6741, 6776), 'click.option', 'click.option', (['"""--connection-string"""'], {}), "('--connection-string')\n", (6753, 6776), False, 'import click\n'), ((6778, 6813), 'click.option', 'click.option', (['"""-o"""', '"""--output-file"""'], {}), "('-o', '--output-file')\n", (6790, 6813), False, 'import click\n'), ((6814, 6841), 'click.option', 'click.option', (['"""--db-schema"""'], {}), "('--db-schema')\n", (6826, 6841), False, 'import click\n'), ((6843, 6877), 'click.option', 'click.option', (['"""--geometry-support"""'], {}), "('--geometry-support')\n", (6855, 6877), False, 'import click\n'), ((6879, 6906), 'click.option', 'click.option', (['"""--from-srid"""'], {}), "('--from-srid')\n", (6891, 6906), False, 'import click\n'), ((6908, 6933), 'click.option', 'click.option', (['"""--to-srid"""'], {}), "('--to-srid')\n", (6920, 6933), False, 'import click\n'), ((6935, 6998), 'click.option', 'click.option', (['"""--logging-config"""'], {'default': '"""logging_config.conf"""'}), "('--logging-config', default='logging_config.conf')\n", (6947, 6998), False, 'import click\n'), ((8175, 8207), 'click.argument', 'click.argument', (['"""new_table_name"""'], {}), "('new_table_name')\n", (8189, 8207), False, 'import click\n'), ((8209, 8241), 'click.argument', 'click.argument', (['"""old_table_name"""'], {}), "('old_table_name')\n", (8223, 8241), False, 'import click\n'), ((8243, 8278), 'click.option', 'click.option', (['"""--connection-string"""'], {}), "('--connection-string')\n", (8255, 8278), False, 'import click\n'), ((8280, 8307), 'click.option', 'click.option', (['"""--db-schema"""'], {}), "('--db-schema')\n", (8292, 8307), False, 'import click\n'), ((8309, 8386), 'click.option', 'click.option', (['"""--select-users"""'], {'help': '"""Users to grant SELECT on updated table"""'}), "('--select-users', help='Users to grant SELECT on updated table')\n", (8321, 8386), False, 'import click\n'), ((8388, 8451), 'click.option', 'click.option', (['"""--logging-config"""'], {'default': '"""logging_config.conf"""'}), "('--logging-config', default='logging_config.conf')\n", (8400, 8451), False, 'import click\n'), ((667, 694), 'logging.getLogger', 'logging.getLogger', (['"""the_el"""'], {}), "('the_el')\n", (684, 694), False, 'import logging\n'), ((1005, 1054), 'os.getenv', 'os.getenv', (['"""CONNECTION_STRING"""', 'connection_string'], {}), "('CONNECTION_STRING', connection_string)\n", (1014, 1054), False, 'import os\n'), ((1347, 1379), 'sqlalchemy.create_engine', 'create_engine', (['connection_string'], {}), '(connection_string)\n', (1360, 1379), False, 'from sqlalchemy import create_engine\n'), ((1394, 1518), 'jsontableschema_sql.Storage', 'Storage', (['engine'], {'dbschema': 'db_schema', 'geometry_support': 'geometry_support', 'from_srid': 'from_srid', 'to_srid': 'to_srid', 'views': '(True)'}), '(engine, dbschema=db_schema, geometry_support=geometry_support,\n    from_srid=from_srid, to_srid=to_srid, views=True)\n', (1401, 1518), False, 'from jsontableschema_sql import Storage\n'), ((9104, 9136), 'sqlalchemy.create_engine', 'create_engine', (['connection_string'], {}), '(connection_string)\n', (9117, 9136), False, 'from sqlalchemy import create_engine\n'), ((467, 485), 'logging.config.dictConfig', 'dictConfig', (['config'], {}), '(config)\n', (477, 485), False, 'from logging.config import dictConfig\n'), ((1725, 1752), 'smart_open.smart_open', 'smart_open', (['file'], {'mode': 'mode'}), '(file, mode=mode)\n', (1735, 1752), False, 'from smart_open import smart_open\n'), ((1975, 1995), 'json.loads', 'json.loads', (['contents'], {}), '(contents)\n', (1985, 1995), False, 'import json\n'), ((2535, 2562), 'json.dump', 'json.dump', (['descriptor', 'file'], {}), '(descriptor, file)\n', (2544, 2562), False, 'import json\n'), ((3360, 3424), 're.match', 're.match', (['carto.carto_connection_string_regex', 'connection_string'], {}), '(carto.carto_connection_string_regex, connection_string)\n', (3368, 3424), False, 'import re\n'), ((5050, 5066), 'csv.reader', 'csv.reader', (['file'], {}), '(file)\n', (5060, 5066), False, 'import csv\n'), ((7526, 7542), 'csv.writer', 'csv.writer', (['file'], {}), '(file)\n', (7536, 7542), False, 'import csv\n'), ((8608, 8672), 're.match', 're.match', (['carto.carto_connection_string_regex', 'connection_string'], {}), '(carto.carto_connection_string_regex, connection_string)\n', (8616, 8672), False, 'import re\n'), ((443, 458), 'yaml.load', 'yaml.load', (['file'], {}), '(file)\n', (452, 458), False, 'import yaml\n'), ((579, 652), 'logging.basicConfig', 'logging.basicConfig', ([], {'format': 'FORMAT', 'level': 'logging.INFO', 'stream': 'sys.stderr'}), '(format=FORMAT, level=logging.INFO, stream=sys.stderr)\n', (598, 652), False, 'import logging\n'), ((5128, 5192), 're.match', 're.match', (['carto.carto_connection_string_regex', 'connection_string'], {}), '(carto.carto_connection_string_regex, connection_string)\n', (5136, 5192), False, 'import re\n'), ((8056, 8073), 'json.dumps', 'json.dumps', (['field'], {}), '(field)\n', (8066, 8073), False, 'import json\n')]