~launchpad-pqm/launchpad/devel

#!/usr/bin/python -S
#
# Copyright 2009 Canonical Ltd.  This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).
#
# This modules uses relative imports.
# pylint: disable-msg=W0403

"""
Add full text indexes to the launchpad database
"""
__metaclass__ = type

import _pythonpath

from distutils.version import LooseVersion
import os.path
from optparse import OptionParser
import subprocess
import sys
from tempfile import NamedTemporaryFile
from textwrap import dedent
import time

import psycopg2.extensions

from canonical import lp
from canonical.database.sqlbase import (
    connect, ISOLATION_LEVEL_AUTOCOMMIT, ISOLATION_LEVEL_READ_COMMITTED,
    quote, quote_identifier)
from canonical.launchpad.scripts import logger, logger_options, db_options

import replication.helpers

# Defines parser and locale to use.
DEFAULT_CONFIG = 'default'

PGSQL_BASE = '/usr/share/postgresql'

A, B, C, D = 'ABCD' # tsearch2 ranking constants

# This data structure defines all of our full text indexes.  Each tuple in the
# top level list creates a 'fti' column in the specified table.
# The letters letters A-D assign a weight to the corresponding column.
# A is most important, and D is least important. This affects result ordering
# when you are ordering by rank.
ALL_FTI = [
    ('archive', [
            ('description', A),
            ('package_description_cache', B),
            ]),
    ('bug', [
            ('name', A),
            ('title', B),
            ('description', D),
            ]),

    ('bugtask', [
            ('targetnamecache', B),
            ('statusexplanation', C),
            ]),

    ('binarypackagerelease', [
            ('summary', B),
            ('description', C),
            ]),

    ('cve', [
            ('sequence', A),
            ('description', B),
            ]),

    ('distribution', [
            ('name', A),
            ('displayname', A),
            ('title', B),
            ('summary', C),
            ('description', D),
            ]),

    ('distributionsourcepackagecache', [
            ('name', A),
            ('binpkgnames', B),
            ('binpkgsummaries', C),
            ('binpkgdescriptions', D),
            ('changelog', D),
            ]),

    ('distroseriespackagecache', [
            ('name', A),
            ('summaries', B),
            ('descriptions', C),
            ]),

    ('faq', [
            ('title', A),
            ('tags', B),
            ('content', D),
            ]),

    ('message', [
            ('subject', B),
            ]),

    ('messagechunk', [
            ('content', C),
            ]),

    ('person', [
            ('name', A),
            ('displayname', A),
            ]),

    ('product', [
            ('name', A),
            ('displayname', A),
            ('title', B),
            ('summary', C),
            ('description', D),
            ]),

    ('productreleasefile', [
            ('description', D),
            ]),

    ('project', [
            ('name', A),
            ('displayname', A),
            ('title', B),
            ('summary', C),
            ('description', D),
            ]),

    ('shippingrequest', [
            ('recipientdisplayname', A),
            ]),

    ('specification', [
            ('name', A),
            ('title', A),
            ('summary', B),
            ('whiteboard', D),
            ]),

    ('question', [
            ('title', A),
            ('description', B),
            ('whiteboard', B),
            ])
    ]


def execute(con, sql, results=False, args=None):
    sql = sql.strip()
    log.debug('* %s' % sql)
    cur = con.cursor()
    if args is None:
        cur.execute(sql)
    else:
        cur.execute(sql, args)
    if results:
        return list(cur.fetchall())
    else:
        return None


def sexecute(con, sql):
    """If we are generating a slonik script, write out the SQL to our
    SQL script. Otherwise execute on the DB.
    """
    if slonik_sql is not None:
        print >> slonik_sql, dedent(sql + ';')
    else:
        execute(con, sql)


def fti(con, table, columns, configuration=DEFAULT_CONFIG):
    """Setup full text indexing for a table"""

    index = "%s_fti" % table
    qindex = quote_identifier(index)
    qtable = quote_identifier(table)
    # Quote the columns
    qcolumns = [
        (quote_identifier(column), weight) for column, weight in columns
        ]

    # Drop the trigger if it exists
    trigger_exists = bool(execute(con, """
        SELECT COUNT(*) FROM pg_trigger, pg_class, pg_namespace
        WHERE pg_trigger.tgname = 'tsvectorupdate'
            AND pg_trigger.tgrelid = pg_class.oid
            AND pg_class.relname = %(table)s
            AND pg_class.relnamespace = pg_namespace.oid
            AND pg_namespace.nspname = 'public'
        """, results=True, args=vars())[0][0])
    if trigger_exists:
        log.debug('tsvectorupdate trigger exists in %s. Dropping.' % qtable)
        sexecute(con, "DROP TRIGGER tsvectorupdate ON %s" % qtable)

    # Drop the fti index if it exists
    index_exists = bool(execute(con, """
        SELECT COUNT(*) FROM pg_index, pg_class, pg_namespace
        WHERE pg_index.indexrelid = pg_class.oid
            AND pg_class.relnamespace = pg_namespace.oid
            AND pg_class.relname = %(index)s
            AND pg_namespace.nspname = 'public'
        """, results=True, args=vars())[0][0])
    if index_exists:
        log.debug('%s exists. Dropping.' % qindex)
        sexecute(con, "DROP INDEX %s" % qindex)

    # Create the 'fti' column if it doesn't already exist
    column_exists = bool(execute(con, """
        SELECT COUNT(*) FROM pg_attribute, pg_class, pg_namespace
        WHERE pg_attribute.attname='fti'
            AND pg_attribute.attisdropped IS FALSE
            AND pg_attribute.attrelid = pg_class.oid
            AND pg_class.relname = %(table)s
            AND pg_class.relnamespace = pg_namespace.oid
            AND pg_namespace.nspname = 'public'
        """, results=True, args=vars())[0][0])
    if not column_exists:
        log.debug('fti column does not exist in %s. Creating.' % qtable)
        sexecute(con, "ALTER TABLE %s ADD COLUMN fti tsvector" % qtable)

    # Create the trigger
    columns_and_weights = []
    for column, weight in qcolumns:
        columns_and_weights.extend( (column, weight) )

    sql = """
        CREATE TRIGGER tsvectorupdate BEFORE UPDATE OR INSERT ON %s
        FOR EACH ROW EXECUTE PROCEDURE ftiupdate(%s)
        """ % (table, ','.join(columns_and_weights))
    sexecute(con, sql)

    # Rebuild the fti column, as the information it contains may be out
    # of date with recent configuration updates.
    sexecute(con, r"""UPDATE %s SET fti=NULL""" % qtable)

    # Create the fti index
    sexecute(con, "CREATE INDEX %s ON %s USING gist(fti)" % (
        qindex, qtable
        ))


def nullify(con):
    """Set all fti index columns to NULL"""
    for table, ignored in ALL_FTI:
        table = quote_identifier(table)
        log.info("Removing full text index data from %s", table)
        sexecute(con, "ALTER TABLE %s DISABLE TRIGGER tsvectorupdate" % table)
        sexecute(con, "UPDATE %s SET fti=NULL" % table)
        sexecute(con, "ALTER TABLE %s ENABLE TRIGGER tsvectorupdate" % table)
    sexecute(con, "DELETE FROM FtiCache")


def liverebuild(con):
    """Rebuild the data in all the fti columns against possibly live database.
    """
    batch_size = 50 # Update maximum of this many rows per commit
    cur = con.cursor()
    for table, ignored in ALL_FTI:
        table = quote_identifier(table)
        cur.execute("SELECT max(id) FROM %s" % table)
        max_id = cur.fetchone()[0]
        if max_id is None:
            log.info("No data in %s - skipping", table)
            continue

        log.info("Rebuilding fti column on %s", table)
        for id in range(0, max_id, batch_size):
            try:
                query = """
                    UPDATE %s SET fti=NULL WHERE id BETWEEN %d AND %d
                    """ % (table, id + 1, id + batch_size)
                log.debug(query)
                cur.execute(query)
            except psycopg2.Error:
                # No commit - we are in autocommit mode
                log.exception('psycopg error')
                con = connect(lp.dbuser)
                con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)


def setup(con, configuration=DEFAULT_CONFIG):
    """Setup and install tsearch2 if isn't already"""

    # tsearch2 is out-of-the-box in 8.3+
    required = LooseVersion('8.3.0')
    assert get_pgversion(con) >= required, (
        'This script only supports PostgreSQL 8.3+')

    schema_exists = bool(execute(
        con, "SELECT COUNT(*) FROM pg_namespace WHERE nspname='ts2'",
        results=True)[0][0])
    if not schema_exists:
        execute(con, 'CREATE SCHEMA ts2')
        con.commit()
    execute(con, 'SET search_path = ts2, public;')

    tsearch2_sql_path = get_tsearch2_sql_path(con)

    ts2_installed = bool(execute(con, """
        SELECT COUNT(*) FROM pg_type,pg_namespace
        WHERE pg_type.typnamespace=pg_namespace.oid
            AND pg_namespace.nspname  = 'ts2'
        """, results=True)[0][0])
    if not ts2_installed:
        assert slonik_sql is None, """
            tsearch2 needs to be setup on each node first with
            fti.py --setup-only
            """

        log.debug('Installing tsearch2')
        cmd = 'psql -f - -d %s' % lp.get_dbname()
        if lp.dbhost:
            cmd += ' -h %s' % lp.dbhost
        if options.dbuser:
            cmd += ' -U %s' % options.dbuser
        p = subprocess.Popen(
            cmd.split(' '), stdin=subprocess.PIPE,
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        out, err = p.communicate(
            "SET client_min_messages=ERROR; CREATE SCHEMA ts2;"
            + open(tsearch2_sql_path).read().replace('public;','ts2, public;'))
        if p.returncode != 0:
            log.fatal('Error executing %s:', cmd)
            log.debug(out)
            sys.exit(rv)

    # Create ftq helper and its sibling _ftq.
    # ftq(text) returns a tsquery, suitable for use querying the full text
    # indexes. _ftq(text) returns the string that would be parsed by
    # to_tsquery and is used to debug the query we generate.
    shared_func = r'''
        import re

        # I think this method would be more robust if we used a real
        # tokenizer and parser to generate the query string, but we need
        # something suitable for use as a stored procedure which currently
        # means no external dependancies.

        # Convert to Unicode
        query = args[0].decode('utf8')
        ## plpy.debug('1 query is %s' % repr(query))

        # Normalize whitespace
        query = re.sub("(?u)\s+"," ", query)

        # Convert AND, OR, NOT and - to tsearch2 punctuation
        query = re.sub(r"(?u)(?:^|\s)-([\w\(])", r" !\1", query)
        query = re.sub(r"(?u)\bAND\b", "&", query)
        query = re.sub(r"(?u)\bOR\b", "|", query)
        query = re.sub(r"(?u)\bNOT\b", " !", query)
        ## plpy.debug('2 query is %s' % repr(query))

        # Deal with unwanted punctuation. We convert strings of punctuation
        # inside words to a '-' character for the hypenation handling below
        # to deal with further. Outside of words we replace with whitespace.
        # We don't mess with -&|!()' as they are handled later.
        #punctuation = re.escape(r'`~@#$%^*+=[]{}:;"<>,.?\/')
        punctuation = r"[^\w\s\-\&\|\!\(\)']"
        query = re.sub(r"(?u)(\w)%s+(\w)" % (punctuation,), r"\1-\2", query)
        query = re.sub(r"(?u)%s+" % (punctuation,), " ", query)
        ## plpy.debug('3 query is %s' % repr(query))

        # Strip ! characters inside and at the end of a word
        query = re.sub(r"(?u)(?<=\w)[\!]+", " ", query)

        # Now that we have handle case sensitive booleans, convert to lowercase
        query = query.lower()

        # Convert foo-bar to ((foo&bar)|foobar) and foo-bar-baz to
        # ((foo&bar&baz)|foobarbaz)
        def hyphen_repl(match):
            bits = match.group(0).split("-")
            return "((%s)|%s)" % ("&".join(bits), "".join(bits))
        query = re.sub(r"(?u)\b\w+-[\w\-]+\b", hyphen_repl, query)
        ## plpy.debug('4 query is %s' % repr(query))

        # Any remaining - characters are spurious
        query = query.replace('-','')

        # Remove unpartnered bracket on the left and right
        query = re.sub(r"(?ux) ^ ( [^(]* ) \)", r"(\1)", query)
        query = re.sub(r"(?ux) \( ( [^)]* ) $", r"(\1)", query)

        # Remove spurious brackets
        query = re.sub(r"(?u)\(([^\&\|]*?)\)", r" \1 ", query)
        ## plpy.debug('5 query is %s' % repr(query))

        # Insert & between tokens without an existing boolean operator
        # ( not proceeded by (|&!
        query = re.sub(r"(?u)(?<![\(\|\&\!])\s*\(", "&(", query)
        ## plpy.debug('6 query is %s' % repr(query))
        # ) not followed by )|&
        query = re.sub(r"(?u)\)(?!\s*(\)|\||\&|\s*$))", ")&", query)
        ## plpy.debug('6.1 query is %s' % repr(query))
        # Whitespace not proceded by (|&! not followed by &|
        query = re.sub(r"(?u)(?<![\(\|\&\!\s])\s+(?![\&\|\s])", "&", query)
        ## plpy.debug('7 query is %s' % repr(query))

        # Detect and repair syntax errors - we are lenient because
        # this input is generally from users.

        # Fix unbalanced brackets
        openings = query.count("(")
        closings = query.count(")")
        if openings > closings:
            query = query + " ) "*(openings-closings)
        elif closings > openings:
            query = " ( "*(closings-openings) + query
        ## plpy.debug('8 query is %s' % repr(query))

        # Strip ' character that do not have letters on both sides
        query = re.sub(r"(?u)((?<!\w)'|'(?!\w))", "", query)

        # Brackets containing nothing but whitespace and booleans, recursive
        last = ""
        while last != query:
            last = query
            query = re.sub(r"(?u)\([\s\&\|\!]*\)", "", query)
        ## plpy.debug('9 query is %s' % repr(query))

        # An & or | following a (
        query = re.sub(r"(?u)(?<=\()[\&\|\s]+", "", query)
        ## plpy.debug('10 query is %s' % repr(query))

        # An &, | or ! immediatly before a )
        query = re.sub(r"(?u)[\&\|\!\s]*[\&\|\!]+\s*(?=\))", "", query)
        ## plpy.debug('11 query is %s' % repr(query))

        # An &,| or ! followed by another boolean.
        query = re.sub(r"(?ux) \s* ( [\&\|\!] ) [\s\&\|]+", r"\1", query)
        ## plpy.debug('12 query is %s' % repr(query))

        # Leading & or |
        query = re.sub(r"(?u)^[\s\&\|]+", "", query)
        ## plpy.debug('13 query is %s' % repr(query))

        # Trailing &, | or !
        query = re.sub(r"(?u)[\&\|\!\s]+$", "", query)
        ## plpy.debug('14 query is %s' % repr(query))

        # If we have nothing but whitespace and tsearch2 operators,
        # return NULL.
        if re.search(r"(?u)^[\&\|\!\s\(\)]*$", query) is not None:
            return None

        # Convert back to UTF-8
        query = query.encode('utf8')
        ## plpy.debug('15 query is %s' % repr(query))
        '''
    text_func = shared_func + """
        return query or None
        """
    tsquery_func = shared_func + """
        p = plpy.prepare("SELECT to_tsquery('%s', $1) AS x", ["text"])
        query = plpy.execute(p, [query], 1)[0]["x"]
        return query or None
        """  % configuration
    sexecute(con, r"""
        CREATE OR REPLACE FUNCTION ts2._ftq(text) RETURNS text AS %s
        LANGUAGE plpythonu IMMUTABLE
        RETURNS NULL ON NULL INPUT
        """ % quote(text_func))
    #print psycopg2.extensions.QuotedString(text_func)
    sexecute(con, r"""
        CREATE OR REPLACE FUNCTION ts2.ftq(text) RETURNS tsquery AS %s
        LANGUAGE plpythonu IMMUTABLE
        RETURNS NULL ON NULL INPUT
        """ % quote(tsquery_func))

    sexecute(con,
            r"COMMENT ON FUNCTION ftq(text) IS '"
            r"Convert a string to a tsearch2 query using the preferred "
            r"configuration. eg. "
            r""""SELECT * FROM Bug WHERE fti @@ ftq(''fatal crash'')". """
            r"The query is lowercased, and multiple words searched using "
            r"AND.'"
            )
    sexecute(con,
            r"COMMENT ON FUNCTION ftq(text) IS '"
            r"Convert a string to an unparsed tsearch2 query'"
            )

    # Create our trigger function. The default one that ships with tsearch2
    # doesn't support weighting so we need our own. We remove safety belts
    # since we know we will be calling it correctly.
    sexecute(con, r"""
        CREATE OR REPLACE FUNCTION ts2.ftiupdate() RETURNS trigger AS '
            new = TD["new"]
            args = TD["args"][:]

            # Short circuit if none of the relevant columns have been
            # modified and fti is not being set to NULL (setting the fti
            # column to NULL is thus how we can force a rebuild of the fti
            # column).
            if TD["event"] == "UPDATE" and new["fti"] != None:
                old = TD["old"]
                relevant_modification = False
                for column_name in args[::2]:
                    if new[column_name] != old[column_name]:
                        relevant_modification = True
                        break
                if not relevant_modification:
                    return "OK"

            # Generate an SQL statement that turns the requested
            # column values into a weighted tsvector
            sql = []
            for i in range(0, len(args), 2):
                sql.append(
                        "ts2.setweight(ts2.to_tsvector(''default'', coalesce("
                        "substring(ltrim($%d) from 1 for 2500),'''')),"
                        "CAST($%d AS \\"char\\"))" % (i + 1, i + 2))
                args[i] = new[args[i]]

            sql = "SELECT %s AS fti" % "||".join(sql)

            # Execute and store in the fti column
            plan = plpy.prepare(sql, ["text", "char"] * (len(args)/2))
            new["fti"] = plpy.execute(plan, args, 1)[0]["fti"]

            # Tell PostgreSQL we have modified the data
            return "MODIFY"
        ' LANGUAGE plpythonu
        """)

    sexecute(con,
        r"COMMENT ON FUNCTION ftiupdate() IS 'Trigger function that keeps "
        r"the fti tsvector column up to date.'"
        )

    # Confirm database locale is valid, and set the 'default' tsearch2
    # configuration to use it.
    r = execute(con, r"""
            SELECT setting FROM pg_settings
            WHERE context='internal' AND name='lc_ctype'
            """, results=True)
    assert len(r) == 1, 'Unable to determine database locale'
    locale = r[0][0]
    assert locale.startswith('en_') or locale in ('C', 'en'), (
            "Non-english database locales are not supported with launchpad. "
            "Fresh initdb required."
            )
    r = locale.split('.', 1)
    if len(r) > 1:
        assert r[1].upper() in ("UTF8", "UTF-8"), \
                "Only UTF8 encodings supported. Fresh initdb required."
    else:
        assert len(r) == 1, 'Invalid database locale %s' % repr(locale)

    r = execute(con,
            "SELECT COUNT(*) FROM pg_ts_config WHERE cfgname='default'",
            results=True)
    if r[0][0] == 0:
        sexecute(con, """
            CREATE TEXT SEARCH CONFIGURATION ts2.default (
                COPY = pg_catalog.english)""")

    # Don't bother with this - the setting is not exported with dumps
    # or propogated  when duplicating the database. Only reliable
    # way we can use is setting search_path in postgresql.conf
    #
    # Set the default schema search path so this stuff can be found
    #execute(con, 'ALTER DATABASE %s SET search_path = public,ts2;' % dbname)


def needs_refresh(con, table, columns):
    '''Return true if the index needs to be rebuilt.

    We know this by looking in our cache to see what the previous
    definitions were, and the --force command line argument
    '''
    current_columns = repr(sorted(columns)) # Convert to a string

    existing = execute(
        con, "SELECT columns FROM FtiCache WHERE tablename=%(table)s",
        results=True, args=vars()
        )
    if len(existing) == 0:
        log.debug("No fticache for %(table)s" % vars())
        sexecute(con, """
            INSERT INTO FtiCache (tablename, columns) VALUES (%s, %s)
            """ % (quote(table), quote(current_columns)))
        return True

    if not options.force:
        previous_columns = existing[0][0]
        if current_columns == previous_columns:
            log.debug("FtiCache for %(table)s still valid" % vars())
            return False
        log.debug("Cache out of date - %s != %s" % (
            current_columns, previous_columns
            ))
    sexecute(con, """
        UPDATE FtiCache SET columns = %s
        WHERE tablename = %s
        """ % (quote(current_columns), quote(table)))

    return True


def get_pgversion(con):
    rows = execute(con, r"show server_version", results=True)
    return LooseVersion(rows[0][0])


def get_tsearch2_sql_path(con):
    major, minor = get_pgversion(con).version[:2]
    path = os.path.join(
        PGSQL_BASE, '%d.%d' % (major, minor), 'contrib', 'tsearch2.sql')
    assert os.path.exists(path), '%s does not exist' % path
    return path


# Script options and arguments parsed from the command line by main()
options = None
args = None

# Logger, setup by main()
log = None

# Files for output generated for slonik(1). None if not a Slony-I install.
slonik_sql = None

def main():
    parser = OptionParser()
    parser.add_option(
            "-s", "--setup-only", dest="setup",
            action="store_true", default=False,
            help="Only install tsearch2 - don't build the indexes.",
            )
    parser.add_option(
            "-f", "--force", dest="force",
            action="store_true", default=False,
            help="Force a rebuild of all full text indexes.",
            )
    parser.add_option(
            "-0", "--null", dest="null",
            action="store_true", default=False,
            help="Set all full text index column values to NULL.",
            )
    parser.add_option(
            "-l", "--live-rebuild", dest="liverebuild",
            action="store_true", default=False,
            help="Rebuild all the indexes against a live database.",
            )
    db_options(parser)
    logger_options(parser)

    global options, args
    (options, args) = parser.parse_args()

    if options.setup + options.force + options.null + options.liverebuild > 1:
        parser.error("Incompatible options")

    global log
    log = logger(options)

    con = connect(lp.dbuser)

    is_replicated_db = replication.helpers.slony_installed(con)

    if options.liverebuild and is_replicated_db:
        parser.error("--live-rebuild does not work with Slony-I install.")

    if is_replicated_db:
        global slonik_sql
        slonik_sql = NamedTemporaryFile(prefix="fti_sl", suffix=".sql")
        print >> slonik_sql, "-- Generated by %s %s" % (
                sys.argv[0], time.ctime())

    if options.liverebuild:
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        liverebuild(con)
    else:
        con.set_isolation_level(ISOLATION_LEVEL_READ_COMMITTED)
        setup(con)
        if options.null:
            nullify(con)
        elif not options.setup:
            for table, columns in ALL_FTI:
                if needs_refresh(con, table, columns):
                    log.info("Rebuilding full text index on %s", table)
                    fti(con, table, columns)
                else:
                    log.info(
                        "No need to rebuild full text index on %s", table)

    if is_replicated_db:
        slonik_sql.flush()
        con.close()
        log.info("Executing generated SQL using slonik")
        if replication.helpers.execute_slonik("""
            execute script (
                set id=@lpmain_set,
                event node=@master_node,
                filename='%s');
            """ % slonik_sql.name, sync=0):
            return 0
        else:
            log.fatal("Failed to execute SQL in Slony-I environment.")
            return 1
    else:
        con.commit()
        return 0


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