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     def __init__(self, msg, table):

    32         Exception.__init__(self, msg)

    33         self.table = table

    34 

    35 

    36 class PgDataCopy:

    37     

    38     def __init__(self, conn1, conn2):

    39         self.conn1 = conn1

    40         self.conn2 = conn2

    41         self.fulltable1 = None

    42         self.fulltable2 = None

    43    

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

    45         self.schema1 = schema

    46         self.table1 = table

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

    48     

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

    50         self.schema2 = schema

    51         self.table2 = table

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

    53     

    54     def copy(self):

    55         self.check()

    56         

    57         buf = io.StringIO()

    58         try:

    59             self.read(buf)

    60             data = buf.getvalue()

    61         finally:

    62             buf.close()

    63         

    64         buf = io.StringIO(data)

    65         try:

    66             self.write(buf)

    67         finally:

    68             buf.close()

    69         

    70         self.analyze()

    71     

    72     def check(self):

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

    74         q = self._compose_check(self.fulltable2)

    75         curs = self.conn2.cursor()

    76         curs.execute(q)

    77         curs.connection.commit()

    78         if curs.rowcount > 0:

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

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

    81     

    82     def read(self, tmpfile):

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

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

    85         curs = self.conn1.cursor()

    86         curs.copy_expert(q, tmpfile)

    87         curs.connection.commit()

    88     

    89     def write(self, tmpfile):

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

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

    92         curs = self.conn2.cursor()

    93         curs.copy_expert(q, tmpfile)

    94         curs.connection.commit()

    95     

    96     def analyze(self):

    97         '''Analyze target table.'''

    98         q = self._compose_analyze(self.fulltable2)

    99         curs = self.conn2.cursor()

   100         curs.execute(q)

   101         curs.connection.commit()

   102     

   103     def _compose_check(self, table):

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

   105     

   106     def _compose_read(self, table, cols):

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

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

   109     

   110     def _compose_write(self, table, cols):

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

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

   113     

   114     def _compose_analyze(self, table):

   115         return 'ANALYZE %s' % table

   116