EPFIngester.py

# Copyright (c) 2010 Apple  Inc. All rights reserved.

# IMPORTANT:  This Apple software is supplied to you by Apple Inc. ("Apple") in
# consideration of your agreement to the following terms, and your use,
# installation, modification or redistribution of this Apple software
# constitutes acceptance of these terms.  If you do not agree with these terms,
# please do not use, install, modify or redistribute this Apple software.

# In consideration of your agreement to abide by the following terms, and subject
# to these terms, Apple grants you a personal, non-exclusive license, under Apple's
# copyrights in this original Apple software (the "Apple Software"), to use,
# reproduce, modify and redistribute the Apple Software, with or without
# modifications, in source and/or binary forms; provided that if you redistribute
# the Apple Software in its entirety and without modifications, you must retain
# this notice and the following text and disclaimers in all such redistributions
# of the Apple Software.  Neither the name, trademarks, service marks or logos of
# Apple Inc. may be used to endorse or promote products derived from the Apple
# Software without specific prior written permission from Apple.  Except as
# expressly stated in this notice, no other rights or licenses, express or implied,
# are granted by Apple herein, including but not limited to any patent rights that
# may be infringed by your derivative works or by other works in which the Apple
# Software may be incorporated.

# The Apple Software is provided by Apple on an "AS IS" basis.  APPLE MAKES NO
# WARRANTIES, EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION THE IMPLIED
# WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE, REGARDING THE APPLE SOFTWARE OR ITS USE AND OPERATION ALONE OR IN
# COMBINATION WITH YOUR PRODUCTS.

# IN NO EVENT SHALL APPLE BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
# GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# ARISING IN ANY WAY OUT OF THE USE, REPRODUCTION, MODIFICATION AND/OR DISTRIBUTION
# OF THE APPLE SOFTWARE, HOWEVER CAUSED AND WHETHER UNDER THEORY OF CONTRACT, TORT
# (INCLUDING NEGLIGENCE), STRICT LIABILITY OR OTHERWISE, EVEN IF APPLE HAS BEEN
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import EPFParser
import MySQLdb
import psycopg2
import os
import datetime
import warnings
import logging

DATETIME_FORMAT = "%y-%m-%d %H:%M:%S"

LOGGER = logging.getLogger()

class Ingester(object):
    """
    Used to ingest an EPF file into a MySQL database.
    """
    #MySQLdb turns MySQL warnings into python warnings, whose behavior is somewhat arcane
    #(as compared with python exceptions.
    #By default, turn all warnings into exceptions
    warnings.filterwarnings('error')
    #Supress warnings that occur when we do a 'DROP TABLE IF EXISTS'; we expect these,
    #so there's no point in cluttering up the output with them.
    warnings.filterwarnings('ignore', 'Unknown table.*')

    def __init__(self,
            filePath,
            tablePrefix=None,
            dbHost='localhost',
            dbUser='epfimporter',
            dbPassword='epf123',
            dbName='epf',
            dbType='mysql',
            recordDelim='\x02\n',
            fieldDelim='\x01'):
        """
        """
        self.filePath = filePath
        self.fileName = os.path.basename(filePath)
        pref = ("%s_" % tablePrefix if tablePrefix else "")
        self.tableName = (pref + self.fileName).replace("-", "_") #hyphens aren't allowed in table names
        self.tableName = self.tableName.split(".")[0]
        self.tmpTableName = self.tableName + "_tmp"
        self.incTableName = self.tableName + "_inc" #used during incremental ingests
        self.unionTableName = self.tableName + "_un" #used during incremental ingests
        self.dbHost = dbHost
        self.dbUser = dbUser
        self.dbPassword = dbPassword
        self.dbName = dbName
        self.dbType = dbType
        self.isPostgresql = (dbType == "postgresql")
        self.isMysql = (dbType == "mysql")
        self.lastRecordIngested = -1
        if self.isPostgresql:
            self.parser = EPFParser.Parser(filePath, typeMap={"VARCHAR(1000)":"TEXT", "VARCHAR(4000)":"TEXT", "CLOB":"LONGTEXT", "DATETIME":"TIMESTAMP", "LONGTEXT":"TEXT"}, recordDelim=recordDelim, fieldDelim=fieldDelim)
        else:
            self.parser = EPFParser.Parser(filePath, recordDelim=recordDelim, fieldDelim=fieldDelim)
        self.startTime = None
        self.endTime = None
        self.abortTime = None
        self.didAbort = False
        self.statusDict = {}
        self.updateStatusDict()
        self.lastRecordCheck = 0
        self.lastTimeCheck = datetime.datetime.now()


    def updateStatusDict(self):
        self.statusDict['fileName'] = self.fileName
        self.statusDict['filePath'] = self.filePath
        self.statusDict['lastRecordIngested'] = self.lastRecordIngested
        self.statusDict['startTime'] = (str(self.startTime) if self.startTime else None)
        self.statusDict['endTime'] = (str(self.endTime) if self.endTime else None)
        self.statusDict['abortTime'] = (str(self.abortTime) if self.abortTime else None)
        self.statusDict['didAbort'] = self.didAbort


    def ingest(self, skipKeyViolators=False):
        """
        Perform a full or incremental ingest, depending on self.parser.exportMode
        """
        if self.parser.exportMode == 'INCREMENTAL':
            self.ingestIncremental(skipKeyViolators=skipKeyViolators)
        else:
            self.ingestFull(skipKeyViolators=skipKeyViolators)


    def ingestFull(self, skipKeyViolators=False):
        """
        Perform a full ingest of the file at self.filePath.

        This is done as follows:
        1. Create a new table with a temporary name
        2. Populate the new table
        3. Drop the old table and rename the new one
        """
        LOGGER.info("Beginning full ingest of %s (%i records)", self.tableName, self.parser.recordsExpected)
        self.startTime = datetime.datetime.now()
        try:
            self._createTable(self.tmpTableName)
            self._populateTable(self.tmpTableName, skipKeyViolators=skipKeyViolators)
            self._renameAndDrop(self.tmpTableName, self.tableName)
        except (MySQLdb.Error, psycopg2.Error), e:
            LOGGER.exception("Fatal error encountered while ingesting '%s'", self.filePath)
            LOGGER.error("Last record ingested before failure: %d", self.lastRecordIngested)
            self.abortTime = datetime.datetime.now()
            self.didAbort = True
            self.updateStatusDict()
            raise #re-raise the exception
        #ingest completed
        self.endTime = datetime.datetime.now()
        self.updateStatusDict()
        LOGGER.info("Full ingest of %s took %s", self.tableName, str(self.endTime - self.startTime))


    def ingestFullResume(self, fromRecord=0, skipKeyViolators=False):
        """
        Resume an interrupted full ingest, continuing from fromRecord.
        """
        LOGGER.info("Resuming full ingest of %s (%i records)", self.tableName, self.parser.recordsExpected)
        self.lastRecordIngested = fromRecord - 1
        self.startTime = datetime.datetime.now()
        try:
            self._populateTable(self.tmpTableName, resumeNum=fromRecord, skipKeyViolators=skipKeyViolators)
            self._renameAndDrop(self.tmpTableName, self.tableName)
        except (MySQLdb.Error, psycopg2.Error), e:
            #LOGGER.error("Error %d: %s", e.args[0], e.args[1])
            LOGGER.error("Error encountered while ingesting '%s'", self.filePath)
            LOGGER.error("Last record ingested before failure: %d", self.lastRecordIngested)
            raise #re-raise the exception
        endTime = datetime.datetime.now()
        ts = str(self.endTime - self.startTime)
        LOGGER.info("Resumed full ingest of %s took %s", self.tableName, ts[:len(ts)-4])


    def ingestIncremental(self, fromRecord=0, skipKeyViolators=False):
        """
        Update the table with the data in the file at filePath.

        If the file to ingest has < 500,000 records, we do a simple REPLACE operation
        on the existing table. If it's larger than that, we use the following 3-step process:
        1. Create a temporary table, and populate it exactly as though it were a Full ingest
        2. Perform a SQL query which selects all rows in the old table whose primary keys *don't*
           match those in the new table, unions the result with all rows in the new table, and
           writes the resulting set to another temporary table.
        3. Swap out the old table for the new one via a rename (same as for Full ingests)
        This proves to be much faster for large files.
        """
        if not (self.tableExists(self.tableName)):
            #The table doesn't exist in the db; this can happen if the full ingest
            #in which the table was added wasn't performed.
            LOGGER.warn("Table '%s' does not exist in the database; skipping", self.tableName)
        else:
            tableColCount = self.columnCount()
            fileColCount = len(self.parser.columnNames)
            assert (tableColCount <= fileColCount) #It's possible for the existing table
            #to have fewer columns than the file we're importing, but it should never have more.

            if fileColCount > tableColCount: #file has "extra" columns
                LOGGER.warn("File contains additional columns not in the existing table. These will not be imported.")
                self.parser.columnNames = self.parser.columnNames[:tableColCount] #trim the columnNames
                # to equal those in the existing table. This will result in the returned records
                # also being sliced.
            s = ("Resuming" if fromRecord else "Beginning")
            LOGGER.info("%s incremental ingest of %s (%i records)", s, self.tableName, self.parser.recordsExpected)
            self.startTime = datetime.datetime.now()

            #Different ingest techniques are faster depending on the size of the input.
            #If there are a large number of records, it's much faster to do a prune-and-merge technique;
            #for fewer records, it's faster to update the existing table.
            try:
                if self.parser.recordsExpected < 500000: #update table in place
                    self._populateTable(self.tableName,
                                    resumeNum=fromRecord,
                                    isIncremental=True,
                                    skipKeyViolators=skipKeyViolators)
                else: #Import as full, then merge the proper records into a new table
                    self._createTable(self.incTableName)
                    LOGGER.info("Populating temporary table...")
                    self._populateTable(self.incTableName, skipKeyViolators=skipKeyViolators)
                    LOGGER.info("Creating merged table...")
                    self._createUnionTable()
                    self._dropTable(self.incTableName)
                    LOGGER.info("Applying primary key constraints...")
                    self._applyPrimaryKeyConstraints(self.unionTableName)
                    self._renameAndDrop(self.unionTableName, self.tableName)

            except (MySQLdb.Error, psycopg2.Error), e:
                #LOGGER.error("Error %d: %s", e.args[0], e.args[1])
                LOGGER.error("Fatal error encountered while ingesting '%s'", self.filePath)
                LOGGER.error("Last record ingested before failure: %d", self.lastRecordIngested)
                self.abortTime = datetime.datetime.now()
                self.didAbort = True
                self.updateStatusDict()
                raise #re-raise the exception
            #ingest completed
            self.endTime = datetime.datetime.now()
            ts = str(self.endTime - self.startTime)
            LOGGER.info("Incremental ingest of %s took %s", self.tableName, ts[:len(ts)-4])
        self.updateStatusDict()


    def connect(self):
        """
        Establish a connection to the database, returning the connection object.
        """
        if self.dbType == "postgresql":
            conn = psycopg2.connect(
                host=self.dbHost,
                user=self.dbUser,
                password=self.dbPassword,
                database=self.dbName)
        else:
            conn = MySQLdb.connect(
                charset='utf8',
                host=self.dbHost,
                user=self.dbUser,
                passwd=self.dbPassword,
                db=self.dbName)

        return conn


    def tableExists(self, tableName=None, connection=None):
        """
        Convenience method which returns True if tableName exists in the db, False if not.

        If tableName is None, uses self.tableName.

        If a connection object is specified, this method uses it and does not close it;
        if not, it creates one using connect(), uses it, and then closes it.
        """

        exStr = """SELECT COUNT(*) FROM information_schema.tables
                    WHERE table_schema = %s
                    AND table_name = %s"""

        tableName = (tableName if tableName else self.tableName)
        conn = (connection if connection else self.connect())
        cur = conn.cursor()
        cur.execute(exStr, (self.dbName, tableName))
        fet = cur.fetchone() #this will always be a 1-tuple; the items's value will be 0 or 1
        doesExist = bool(fet[0])
        cur.close()
        if not connection:
            conn.close()
        return doesExist


    def columnCount(self, tableName=None, connection=None):
        """
        Convenience method for returning the number of columns in tableName.

        If tableName is None, uses self.tableName.

        If a connection object is specified, this method uses it and does not close it;
        if not, it creates one using connect(), uses it, and then closes it.
        """
        tableName = (tableName if tableName else self.tableName)
        conn = (connection if connection else self.connect())
        cur = conn.cursor()
        exStr = """SHOW COLUMNS FROM %s""" % tableName
        colCount = cur.execute(exStr) #cur.execute() returns the number of rows,
        # which for SHOW COLUMNS is the number of columns in the table
        cur.close()
        if not connection:
            conn.close()
        return colCount


    def _createTable(self, tableName):
        """
        Connect to the db and create a table named self.tableName_TMP, dropping previous one if it exists.

        Also adds primary key constraint to the new table.
        """
        conn = self.connect()
        cur = conn.cursor()
        cur.execute("""DROP TABLE IF EXISTS %s""" % tableName)
        #create the column name part of the table creation string
        colPairs = zip(self.parser.columnNames, self.parser.dataTypes)
        lst = [" ".join(aPair) for aPair in colPairs] #list comprehension
        paramStr = ",".join(lst)
        #paramString now looks like "export_date BIGINT, storefront_id INT, country_code VARCHAR(100)" etc.
        exStr = """CREATE TABLE %s (%s)""" % (tableName, paramStr)
        cur.execute(exStr) #create the table in the database
        #set the primary key
        if self.isPostgresql:
            conn.commit()
        conn.close()
        self._applyPrimaryKeyConstraints(tableName)


    def _applyPrimaryKeyConstraints(self, tableName):
        """
        Apply the primary key specified in parser to tableName.
        """
        pkLst = self.parser.primaryKey

        if pkLst:
            conn = self.connect()
            cur = conn.cursor()
            pkStr = ", ".join(pkLst)
            if self.isPostgresql:
                exStr = """ALTER TABLE %s ADD CONSTRAINT %s_pk PRIMARY KEY (%s)""" % (tableName, tableName, pkStr)
            else:
                exStr = """ALTER TABLE %s ADD CONSTRAINT PRIMARY KEY (%s)""" % (tableName, pkStr)
            cur.execute(exStr)
            if self.isPostgresql:
                conn.commit()
            conn.close()


    def _escapeRecords(self, recordList, connection=None):
        """
        Appropriately escape the contents of a list of records (as returned by the parser)
        so that there are no illegal characters (e.g. internal quotes) in the SQL query.

        This is done here rather than in the parser because it uses the literal() method of the
        connection object.
        """
        conn = (connection if connection else self.connect())
        escapedRecords = []
        cur = conn.cursor()
        keys = {}
        for aRec in recordList:
            marker = tuple([aRec[i] for i in self.parser.primaryKeyIndexes])
            if marker in keys: continue
            keys[marker] = 1
            if self.isMysql:
                escRec = [conn.literal(aField) for aField in aRec]
            else:
                escRec = [cur.mogrify("%s", (aField,)) for aField in aRec]
            escapedRecords.append(escRec)
        return escapedRecords


    def _populateTable(self, tableName, resumeNum=0, isIncremental=False, skipKeyViolators=False):
        """
        Populate tableName with data fetched by the parser, first advancing to resumePos.

        For Full imports, if skipKeyViolators is True, any insertions which would violate the primary key constraint
        will be skipped and won't log errors.
        """
        #REPLACE is a MySQL extension which inserts if the key is new, or deletes and inserts if the key is a duplicate
        commandString = ("REPLACE" if (isIncremental and self.isMysql) else "INSERT")
        ignoreString = ("IGNORE" if (skipKeyViolators and not isIncremental and self.isMysql) else "")
        exStrTemplate = """%s %s INTO %s %s VALUES %s"""
        colNamesStr = "(%s)" % (", ".join(self.parser.columnNames))

        self.parser.seekToRecord(resumeNum) #advance to resumeNum
        conn = self.connect()
        if self.isPostgresql:
            cur = conn.cursor()
            createRuleTemplate = """CREATE OR REPLACE RULE %s_on_duplicate_ignore AS ON INSERT TO %s WHERE EXISTS(SELECT 1 FROM %s WHERE (%s) = (%s)) DO INSTEAD NOTHING"""
            pkLst = self.parser.primaryKey
            if len(pkLst) > 0:
                pk1 = ", ".join(pkLst)
                pk2 = ", ".join(["NEW.%s" % pk for pk in pkLst])
            else:
                pk1 = "1"
                pk2 = "1"

            exStr = createRuleTemplate % (tableName, tableName, tableName, pk1, pk2)
            cur.execute(exStr)

        while (True):
            #By default, we concatenate 200 inserts into a single INSERT statement.
            #a large batch size per insert improves performance, until you start hitting max_packet_size issues.
            #If you increase MySQL server's max_packet_size, you may get increased performance by increasing maxNum
            records = self.parser.nextRecords(maxNum=200)
            if (not records):
                break

            escapedRecords = self._escapeRecords(records) #This will sanitize the records
            stringList = ["(%s)" % (", ".join(aRecord)) for aRecord in escapedRecords]

            cur = conn.cursor()
            colVals = unicode(", ".join(stringList), 'utf-8')
            exStr = exStrTemplate % (commandString, ignoreString, tableName, colNamesStr, colVals)
            #unquote NULLs
            exStr = exStr.replace("'NULL'", "NULL")
            exStr = exStr.replace("'null'", "NULL")

            try:
                cur.execute(exStr)
            except (MySQLdb.Warning, psycopg2.Warning), e:
                LOGGER.warning(str(e))
            except (MySQLdb.IntegrityError, psycopg2.IntegrityError), e:
            #This is likely a primary key constraint violation; should only be hit if skipKeyViolators is False
                LOGGER.error(str(e))
            except (MySQLdb.Error, psycopg2.Error), e:
                LOGGER.error("error executing %s" % exStr)
                raise #re-raise the exception
            self.lastRecordIngested = self.parser.latestRecordNum
            recCheck = self._checkProgress()
            if recCheck:
                LOGGER.info("...at record %i...", recCheck)

        if self.isPostgresql:
            dropRuleTemplate = """DROP RULE %s_on_duplicate_ignore ON %s"""
            exStr = dropRuleTemplate % (tableName, tableName)
            cur.execute(exStr)

        if self.isPostgresql:
            conn.commit()

        conn.close()


    def _checkProgress(self, recordGap=5000, timeGap=datetime.timedelta(0, 120, 0)):
        """
        Checks whether recordGap or more records have been ingested since the last check;
        if so, checks whether timeGap seconds have elapsed since the last check.

        If both checks pass, returns self.lastRecordIngested; otherwise returns None.
        """
        if self.lastRecordIngested - self.lastRecordCheck >= recordGap:
            t = datetime.datetime.now()
            if t - self.lastTimeCheck >= timeGap:
                self.lastTimeCheck = t
                self.lastRecordCheck = self.lastRecordIngested
                return self.lastRecordCheck
        return None


    def _dropTable(self, tableName):
        """A convenience method that just connects, drops tableName if it exists, and disconnects"""
        conn = self.connect()
        cur = conn.cursor()
        cur.execute("""DROP TABLE IF EXISTS %s""" % tableName)
        if self.isPostgresql:
            conn.commit()
        conn.close()


    def _renameAndDrop(self, sourceTable, targetTable):
        """
        Temporarily rename targetTable, then rename sourceTable to targetTable.
        If this succeeds, drop the renamed targetTable; otherwise revert it and drop sourceTable.
        """
        conn = self.connect()
        cur = conn.cursor()
        revert = False
        #first, rename the existing "real" table, so we can restore it if something goes wrong
        targetOld = targetTable + "_old"
        cur.execute("""DROP TABLE IF EXISTS %s""" % targetOld)
        if self.isMysql:
            exStr = """ALTER %s %s RENAME %s"""
        else:
            exStr = """ALTER %s %s RENAME TO %s"""

        if (self.tableExists(targetTable, connection=conn)):
            cur.execute(exStr % ("TABLE", targetTable, targetOld))
            if self.isPostgresql:
                cur.execute(exStr % ("INDEX", targetTable+'_pk', targetOld+'_pk'))
        #now rename the new table to replace the old table
        try:
            cur.execute(exStr % ("TABLE", sourceTable, targetTable))
            if self.isPostgresql:
                cur.execute(exStr % ("INDEX", sourceTable+'_pk', targetTable+'_pk'))
        except MySQLdb.Error, e:
            LOGGER.error("Error %d: %s", e.args[0], e.args[1])
            revert = True
        except psycopg2.Error, e:
            LOGGER.error("Error %s", e)
            revert = True

        if revert:
            LOGGER.error("Could not rename tmp table; reverting to original table (if it exists).")
            if (self.tableExists(targetOld, connection=conn)):
                cur.execute(exStr % ("TABLE", targetOld, targetTable))
                if self.isPostgresql:
                    cur.execute(exStr % ("INDEX", targetOld+'_pk', targetTable+'_pk'))
        #Drop sourceTable so it's not hanging around
        #drop the old table
        cur.execute("""DROP TABLE IF EXISTS %s""" % targetOld)
        if self.isPostgresql:
            conn.commit()
        conn.close()


    def _createUnionTable(self):
        """
        After incremental ingest data has been written to self.incTableName, union the pruned
        original table and the new table into a tmp table
        """
        conn = self.connect()
        cur = conn.cursor()
        cur.execute("""DROP TABLE IF EXISTS %s""" % self.unionTableName)
        exStr = """CREATE TABLE %s %s""" % (self.unionTableName, self._incrementalUnionString())
        cur.execute(exStr)
        if self.isPostgresql:
            conn.commit()
        conn.close()


    def _incrementalWhereClause(self):
        """
        Creates and returns the appropriate WHERE clause string used when pruning the target table
        during an incremental ingest
        """
        pCols = self.parser.primaryKey
        substrings = ["%s.%s=%s.%s" % (self.tableName, aCol, self.incTableName, aCol) for aCol in pCols]
        joinedString = " AND ".join(substrings)
        whereClause = "WHERE %s.export_date <= %s.export_date AND %s" % (self.tableName, self.incTableName, joinedString)
        return whereClause


    def _incrementalSelectString(self):
        """
        Creates and returns the appropriate SELECT statement used when pruning the target table
        during an incremental ingest
        """
        whereClause = self._incrementalWhereClause()
        selectString = ("SELECT * FROM %s WHERE 0 = (SELECT COUNT(*) FROM %s %s)" %
            (self.tableName, self.incTableName, whereClause))
        return selectString


    def _incrementalUnionString(self):
        """
        Creates and returns the appropriate UNION string used when merging the pruned table
        with the temporary incrmental table.

        The ingest and pruning process should preclude any dupes, so we can use ALL, which should be faster.
        """
        selectString = self._incrementalSelectString()
        unionString = "IGNORE SELECT * FROM %s UNION ALL %s" % (self.incTableName, selectString)
        return unionString