find_bad_toast

#!/usr/bin/env python

from __future__ import print_function # You must use Python >= 2.6!

import argparse
import os
import psycopg2

def debug(*args):
    if os.getenv('FIND_BAD_TOAST_DEBUG', False):
        print(*args)


"""
find_bad_toast - *****************

Scan a PostgreSQL table for rows with corrupted TOAST records.

Arguments:

- connect_string is psycopg2 connection string.

- table is the name of the table to scan.

- pks is a sequence of column names. These are the columns that compose
the primary key of the column. Of course, there may be only one such column.

This is a translation, for performance reasons, from plpgsql to Python,
of the logic described at
*** http://www.databasesoup.com/2014/07/improved-toast-corruption-function.html

The performance reasons are a) to avoid keeping a transaction open for the
initial whole-table scan, which is painful for large tables, and b) to improve
upon the performance of plpgsql as a language.

"""
def find_bad_toast(connect_string, table, pks):
    conn = psycopg2.connect(connect_string)
    conn.autocommit = True
    
    ids_cur = conn.cursor()
    pk_cols = ', '.join(pks)
    ids_query = 'SELECT {p} FROM {table}'.format(p=pk_cols, table=table)
    debug('Selecting IDs using query "{ids_query}"'.format(ids_query=ids_query))
    ids_cur.execute(ids_query)

    count = 0
    copy_cur = conn.cursor()
    while True:
        ids = ids_cur.fetchone()
        if not ids:
            break

        debug('Examining row with IDs', ids)

        count += 1
        if count % 100000 == 0:
            print('{count} rows inspected'.format(count=count))

        try:
            copy_query = 'COPY (SELECT * FROM {t} WHERE '.format(t=table)

            where_clauses = []
            for col in pks:
                clause = col + ' = %s'
                debug('Adding clause "{clause}"'.format(clause=clause))
                where_clauses.append(clause)

            debug('AND_clauses: ', where_clauses)
            copy_query += ' AND '.join(where_clauses)

            copy_query += ") TO '/tmp/testout'"    

            debug('COPY query:', copy_query)
            copy_cur.execute(copy_query, ids)

        except Exception as e:
            print('Found corruption in the row with PK columns', ids,
                  "Exception: '{ex}'".format(ex=str(e)))


def main():
    parser = argparse.ArgumentParser(
        description='Find rows in a table that have corrupted TOAST values.'
    )
    parser.add_argument('-t', '--table', required=True)
    parser.add_argument('-p', '--pk', required=True, action='append')
    parser.add_argument('--connect-string', required=True)
    args = parser.parse_args()
    
    find_bad_toast(args.connect_string, args.table, args.pk)


if __name__ == '__main__':
    main()