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

     2 #

     3 # PgDataCopy - copy data between tables

     4 #

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

     6 #

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

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

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

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

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

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

    13 #

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

    15 # all copies or substantial portions of the Software.

    16 #

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

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

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

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

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

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

    23 # THE SOFTWARE.

    24 

    25 

    26 import io

    27 

    28 

    29 class TargetNotEmptyError(Exception):

    30     

    31     pass

    32 

    33 

    34 class PgDataCopy:

    35     

    36     def __init__(self, conn1, conn2):

    37         self.conn1 = conn1

    38         self.conn2 = conn2

    39         self.fulltable1 = None

    40         self.fulltable2 = None

    41    

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

    43         self.schema1 = schema

    44         self.table1 = table

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

    46     

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

    48         self.schema2 = schema

    49         self.table2 = table

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

    51     

    52     def copy(self):

    53         self.check()

    54         

    55         buf = io.StringIO()

    56         try:

    57             self.read(buf)

    58             data = buf.getvalue()

    59         finally:

    60             buf.close()

    61         

    62         buf = io.StringIO(data)

    63         try:

    64             self.write(buf)

    65         finally:

    66             buf.close()

    67         

    68         self.analyze()

    69     

    70     def check(self):

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

    72         q = self._compose_check(self.fulltable2)

    73         curs = self.conn2.cursor()

    74         curs.execute(q)

    75         curs.connection.commit()

    76         if curs.rowcount > 0:

    77             raise TargetNotEmptyError('Target table contains data: %s' % self.fulltable2)

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

    79     

    80     def read(self, tmpfile):

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

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

    83         curs = self.conn1.cursor()

    84         curs.copy_expert(q, tmpfile)

    85         curs.connection.commit()

    86     

    87     def write(self, tmpfile):

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

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

    90         curs = self.conn2.cursor()

    91         curs.copy_expert(q, tmpfile)

    92         curs.connection.commit()

    93     

    94     def analyze(self):

    95         '''Analyze target table.'''

    96         q = self._compose_analyze(self.fulltable2)

    97         curs = self.conn2.cursor()

    98         curs.execute(q)

    99         curs.connection.commit()

   100     

   101     def _compose_check(self, table):

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

   103     

   104     def _compose_read(self, table, cols):

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

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

   107     

   108     def _compose_write(self, table, cols):

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

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

   111     

   112     def _compose_analyze(self, table):

   113         return 'ANALYZE %s' % table

   114