root / pykota / trunk / initscripts / postgresql / upgrade-from-before-1.03.py @ 3501

#! /usr/bin/env python
#
# PyKota - Print Quotas for CUPS
#
# (c) 2003-2009 Jerome Alet <alet@librelogiciel.com>
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# $Id$
#
# THIS IS AN UPGRADE SCRIPT FOR OLD PYKOTA DATABASES TO THE 1.03 VERSION
#

import sys
import os
try :
    import pg
except ImportError :
    sys.stderr.write("The PyGreSQL Python module doesn't seem to be available. ABORTED.\n")
    sys.exit(-1)

def dump_old_database() :
    """Dumps the existing PyKota database to a file, to avoir loosing data.

       Returns 1 if dump is successfull, 0 if it isn't.
    """
    pipeinput = os.popen("pg_dump -C -D -N -U postgres -f pykota-dump.sql pykota")
    dummy = pipeinput.read()
    return (pipeinput.close() is None) or 0

def create_new_database() :
    """Creates the empty database."""
    pipeinput = os.popen("psql -U postgres template1 <pykota-postgresql.sql")
    dummy = pipeinput.read()
    return (pipeinput.close() is None) or 0

def restore_original_database() :
    """Creates the empty database."""
    pipeinput = os.popen("psql -U postgres template1 <pykota-dump.sql")
    dummy = pipeinput.read()
    return (pipeinput.close() is None) or 0

def open_database(dbname="pykota") :
    """Returns the database object or None if we can't connect to it."""
    try :
        pykotadb = pg.connect(host="localhost", port=5432, dbname=dbname, user="postgres")
    except pg.error, msg :
        sys.stderr.write("%s\n" % msg)
        sys.stderr.write("Unable to connect to the local PostgreSQL server.\nPlease modify the open_database() method in %s\nto connect to the correct PostgreSQL server\nand relaunch the script.\n" % sys.argv[0])
        return
    else :
        return pykotadb


def doQuote(field) :
    """Quotes a field for use as a string in SQL queries."""
    if type(field) == type(0) : # TODO : do something safer
        typ = "decimal"
    else :
        typ = "text"
    return pg._quote(field, typ)

def main() :
    """Does the work."""

    # First we make a dump of the old database
    print "Dumping old database for safety...",
    if not dump_old_database() :
        sys.stderr.write("Error while dumping old database. ABORTED.\n")
        return -1
    print "Done."

    # Second we try to connect to it
    print "Extracting datas from old database...",
    db = open_database()
    if db is None :
        sys.stderr.write("Impossible to connect to old PyKota database. ABORTED.\nAre you sure you are upgrading an existing installation ?\n")
        return -1

    # Third we extract datas
    oldprinters = db.query("SELECT * FROM printers ORDER BY id;")
    oldusers = db.query("SELECT * FROM users ORDER BY id;")
    oldquotas = db.query("SELECT * FROM userpquota ORDER BY printerid, userid;")

    # Fourth close the database
    db.close()
    print "Done."

    # Fifth we delete the old database !
    answer = raw_input("The old database will be deleted for the upgrade to take place.\nAre you sure you want to continue (y/N) ? ")
    if answer[0:1].upper() != 'Y' :
        sys.stderr.write("User wants to stop now ! ABORTED.\n")
        return -1
    db = open_database("template1")
    if db is None :
        sys.stderr.write("Impossible to connect to database template1. ABORTED.\nPlease report to your Database administrator.\n")
        return -1
    try :
        db.query("DROP DATABASE pykota;")
    except pg.error, msg :
        sys.stderr.write("Impossible to delete old database. ABORTED.\n%s\n" % msg)
        db.close()
        return -1
    else :
        db.close()

    # Sixth we create the new database
    print "Creating the new database...",
    if not create_new_database() :
        sys.stderr.write("impossible to create new database ! ABORTED.\n")
        return -1
    print "Done."

    # Seventh we restore old data
    print "Restoring old datas..."
    db = open_database()
    if db is None :
        sys.stderr.write("Impossible to connect to new PyKota database. ABORTED.\nOld database will be restored if possible.")
        print "An error occured, I'll try to restore the original database..."
        if not restore_original_database() :
            sys.stderr.write("Shit ! A double-error occured !!!\nPlease report problem to your database administrator.\n")
            sys.stderr.write("And file a bug report to alet@librelogiciel.com\n")
        else :
            print "Done."
        return -1
    db.query("BEGIN;")
    try :
        newprinters = {}
        for oldprinter in oldprinters.dictresult() :
            db.query("INSERT INTO printers (printername) VALUES (%s);" % doQuote(oldprinter["printername"]))
            newid = db.query("SELECT id FROM printers WHERE printername='%s';" % oldprinter["printername"]).dictresult()[0]["id"]
            newprinters[oldprinter["id"]] = newid
        newusers = {}
        for olduser in oldusers.dictresult() :
            db.query("INSERT INTO users (username) VALUES (%s);" % doQuote(olduser["username"]))
            newid = db.query("SELECT id FROM users WHERE username='%s';" % olduser["username"]).dictresult()[0]["id"]
            newusers[olduser["id"]] = newid
        for oldquota in oldquotas.dictresult() :
            db.query("INSERT INTO userpquota (userid, printerid, pagecounter, lifepagecounter, softlimit, hardlimit, datelimit) VALUES (%s, %s, %s, %s, %s, %s, %s);" %
                                              (doQuote(newusers[oldquota["userid"]]) , doQuote(newprinters[oldquota["printerid"]]), doQuote(oldquota["pagecounter"]), doQuote(oldquota["lifepagecounter"]), doQuote(oldquota["softlimit"]), doQuote(oldquota["hardlimit"]), doQuote(oldquota["datelimit"])))
    except pg.error, msg :
        sys.stderr.write("ERROR : %s\nABORTED.\n" % msg)
        db.query("ROLLBACK;")
        db.close()
        return -1
    except :
        sys.stderr.write("Unknown error ! ABORTED.\n")
        db.query("ROLLBACK;")
        db.close()
        return -1
    else :
        db.query("COMMIT;")
        db.close()
    print "Done."
    print "NB : Last job on each printer was lost. This is normal !"
    return 0

if __name__ == "__main__" :
    sys.exit(main())