diff -r bd0beda49bcb -r 98c7809af415 pgtoolkit/pgdatacopy.py --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/pgtoolkit/pgdatacopy.py Thu May 10 08:42:21 2012 +0200 @@ -0,0 +1,114 @@ +# -*- coding: utf-8 -*- +# +# PgDataCopy - copy data between tables +# +# Copyright (c) 2012 Radek Brich +# +# 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): + + pass + + +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: %s' % 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 +