#!/usr/bin/env python

""" this script is an on-line alternative to VACUUM FULL CONCURRENTLY;

it uses non-functional updates to move data
in postgreSQL table to lower pages so that VACUUM
can free the now empty the pages at the end of table.

Each run of this script will move ROWS_TO_MOVE pages

The script may run slower on tables where there are Foreign Keys 
pointing to it. Make sure that the FK side columns are indexed too!

NB! - this script does not run VACUUM before start, 
so run VACUUM once before it to make sure there is free space in lower tables

On actively used tables try to move just a portion of rows so you don't fill 
up all pages and cause the main workload to create new pages at the very end 

"""

import sys, pprint
import psycopg2

# configutration here
database = "dbname=sourcebox port=5435"
table = "shrinktest"
pkcol = "id"
ROWS_TO_MOVE = 100000
# configuration end

tablesize_q = "select pg_relation_size(%s)" % (table,)

last_live_rows_q = "SELECT ctid, %s from %s order by 1 desc limit %s" % (pkcol, table, ROWS_TO_MOVE)

con = psycopg2.connect(database)
con.autocommit = False
cur = con.cursor()

def get_current_table_size(tablename):
    cur.execute("select pg_relation_size(%s)", [tablename])
    return cur.fetchone()[0]

startsize = get_current_table_size(table)

print 'startsize:', startsize

def get_last_live_row_ids():
    cur.execute(last_live_rows_q)
    rows = cur.fetchall()
    return rows
    
rows = get_last_live_row_ids()

pprint.pprint(rows[:5])

move_item_to_other_page_q = "UPDATE %s SET %s=%s WHERE %s=%%s RETURNING ctid" % (table, pkcol, pkcol, pkcol)

def move_item_to_other_page(item_row):
    count = 0
    ctid, row_id = item_row
    row_page, row_nr = eval(ctid)
#    print 'id: %s, ctid: (%s,%s)' % ( row_id, row_page, row_nr), 
    while 1:
        count += 1
        cur.execute(move_item_to_other_page_q, [row_id])
        res = cur.fetchone()[0]
        new_page, new_nr = eval(res)
        if new_page <> row_page:
            break
#    print '--> (%s, %s), updates: %s' % (new_page, new_nr, count)
    return new_page

last_page = eval(rows[0][0])[0]

for row in rows:
    page = eval(row[0])[0]
    if page != last_page:
        con.commit()
        last_page = page
        print 'COMMIT. Cleaned page %s (moved to page %s), last live row was at %s' % (last_page,new_page,last_page * 8192)
    new_page = move_item_to_other_page(row)
    if new_page > page:
        con.rollback()
        print 'MOVED TO HIGHER PAGE - NO FREE SPACE TO MOVE LOWER'
        print 'LOWER "ROWS_TO_MOVE" AND TRY AGAIN'
        con.autocommit = True       
        sys.exit()

if new_page < page:
    con.commit()
    last_page = page
    print 'MOVING %s LAST ROWS DONE!' % ROWS_TO_MOVE

con.autocommit = True

print 'RUNNING VACUUM %s' % (table,)
cur.execute('VACUUM %s' % (table,))
print 'DONE!'

endsize = get_current_table_size(table)

print 'startsize:', startsize
print 'endsize:', endsize
print 'shrinked by %s bytes' % (startsize - endsize)