pgtoolkit/pgdatacopy.py
changeset 104 d8ff52a0390f
parent 103 24e94a3da209
child 105 10551741f61f
equal deleted inserted replaced
103:24e94a3da209 104:d8ff52a0390f
     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