# -*- coding: utf-8 -*-

#

# PgDataCopy - copy data between tables

#

# Copyright (c) 2012  Radek Brich <radek.brich@devl.cz>

#

# Permission is hereby granted, free of charge, to any person obtaining a copy

# of this software and associated documentation files (the "Software"), to deal

# in the Software without restriction, including without limitation the rights

# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

# copies of the Software, and to permit persons to whom the Software is

# furnished to do so, subject to the following conditions:

#

# The above copyright notice and this permission notice shall be included in

# all copies or substantial portions of the Software.

#

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

# THE SOFTWARE.

import io

class TargetNotEmptyError(Exception):

    def __init__(self, msg, table):

        Exception.__init__(self, msg)

        self.table = table

class PgDataCopy:

    def __init__(self, conn1, conn2):

        self.conn1 = conn1

        self.conn2 = conn2

        self.fulltable1 = None

        self.fulltable2 = None

    def set_source(self, table, schema='public'):

        self.schema1 = schema

        self.table1 = table

        self.fulltable1 = '"' + schema + '"."'+ table + '"'

    def set_destination(self, table, schema='public'):

        self.schema2 = schema

        self.table2 = table

        self.fulltable2 = '"' + schema + '"."'+ table + '"'

    def copy(self):

        self.check()

        buf = io.StringIO()

        try:

            self.read(buf)

            data = buf.getvalue()

        finally:

            buf.close()

        buf = io.StringIO(data)

        try:

            self.write(buf)

        finally:

            buf.close()

        self.analyze()

    def check(self):

        '''Check that target table does not contain any data (otherwise cannot copy).'''

        q = self._compose_check(self.fulltable2)

        curs = self.conn2.cursor()

        curs.execute(q)

        curs.connection.commit()

        if curs.rowcount > 0:

            raise TargetNotEmptyError('Target table contains data.', self.fulltable2)

        self.cols = [desc[0] for desc in curs.description]

    def read(self, tmpfile):

        '''Read contents from source table.'''

        q = self._compose_read(self.fulltable1, self.cols)

        curs = self.conn1.cursor()

        curs.copy_expert(q, tmpfile)

        curs.connection.commit()

    def write(self, tmpfile):

        '''Write source table contents to target table.'''

        q = self._compose_write(self.fulltable2, self.cols)

        curs = self.conn2.cursor()

        curs.copy_expert(q, tmpfile)

        curs.connection.commit()

    def analyze(self):

        '''Analyze target table.'''

        q = self._compose_analyze(self.fulltable2)

        curs = self.conn2.cursor()

        curs.execute(q)

        curs.connection.commit()

    def _compose_check(self, table):

        return 'SELECT * FROM %s LIMIT 1' % table

    def _compose_read(self, table, cols):

        collist = ', '.join(['"%s"' % col for col in cols])

        return 'COPY %s (%s) TO STDOUT' % (table, collist)

    def _compose_write(self, table, cols):

        collist = ', '.join(['"%s"' % col for col in cols])

        return 'COPY %s (%s) FROM STDIN' % (table, collist)

    def _compose_analyze(self, table):

        return 'ANALYZE %s' % table