Populate wp10 tables

populate_wp10.py

"""
Populate wp10 table based on the generated dump file

See T106278 and T135684

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, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
http://www.gnu.org/copyleft/gpl.html
"""
import argparse
import sys
import codecs
import MySQLdb


def main():
    args = sys.argv[1:]
    parser = argparse.ArgumentParser(description='Parse and add statements '
                                     'based on trained model')
    parser.add_argument('--wiki', '-w', nargs='?', required=True,
                        help='Wiki db name (e.g. enwiki)')
    parser.add_argument('--date', '-d', nargs='?', required=True,
                        help='Dump date e.g. 20160820')
    parser.add_argument('--file', '-f', nargs='?', required=True,
                        help='Path to file')
    parsed_args = parser.parse_known_args(args)[0]

    wiki = parsed_args.wiki
    month = parsed_args.date[:6]
    data_file = codecs.open(parsed_args.file, 'r', 'utf-8')

    db = MySQLdb.connect(host="tools-db", db="s52709__wp10_p",
                         read_default_file="~/replica.my.cnf")
    cursor = db.cursor()
    table_query = """-- ORES article quality scores for enwiki, dump 20160820
--
CREATE TABLE {wiki}_{date} (
    -- Primary key ID
    id INTEGER(10) unsigned NOT NULL PRIMARY KEY AUTO_INCREMENT,
    -- Page id (foreign key to page.page_id)
    page_id INTEGER(10) unsigned NOT NULL,
    -- Title
    title VARCHAR(300) NOT NULL,
    -- Rev id (foreign key to revision.rev_id)
    rev_id INTEGER(10) NOT NULL,
    -- predicted classification
    prediction TINYINT(1) NOT NULL,
    -- Weighted sum, a measure for quality, see T135684
    weighted_sum DECIMAL(3,3) NOT NULL
)

CREATE INDEX quality ON {wiki}_{date} (weighted_sum, prediction, title);
CREATE INDEX class ON {wiki}_{date} (prediction, weighted_sum, title);
CREATE UNIQUE INDEX page ON {wiki}_{date} (title);
    """.format(wiki=wiki, date=month)
    try:
        cursor.execute(table_query)
    except:
        pass
    db.commit()
    cursor.close()
    db.close()

    db = MySQLdb.connect(host="tools-db", db="s51100__wp10_p",
                         read_default_file="~/replica.my.cnf")
    cursor = db.cursor()
    insert_statement = (
        "INSERT INTO {wiki}_{date} "
        "(page_id, title, rev_id, predication, weighted_sum) "
        "VALUES (%s, %s, %s, %s, %s)".format(wiki=wiki, date=month))

    with data_file as f:
        first = True
        for line in f:
            line = line.replace('\n', '')
            if u'\t' not in line:
                continue
            line = line.strip().split('\t')
            if not first:
                cursor.execute(insert_statement,
                               tuple(line))
            first = False
    db.commit()
    cursor.close()
    db.close()


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