/apiary/mysql/mysql.py

  1#
  2# $LicenseInfo:firstyear=2010&license=mit$
  3# 
  4# Copyright (c) 2010, Linden Research, Inc.
  5# 
  6# Permission is hereby granted, free of charge, to any person obtaining a copy
  7# of this software and associated documentation files (the "Software"), to deal
  8# in the Software without restriction, including without limitation the rights
  9# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 10# copies of the Software, and to permit persons to whom the Software is
 11# furnished to do so, subject to the following conditions:
 12# 
 13# The above copyright notice and this permission notice shall be included in
 14# all copies or substantial portions of the Software.
 15# 
 16# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 17# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 18# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 19# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 20# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 21# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 22# THE SOFTWARE.
 23# $/LicenseInfo$
 24#
 25
 26from optparse import OptionParser
 27
 28import MySQLdb
 29import os
 30import re
 31import sys
 32import time
 33
 34import apiary
 35from apiary.tools.debug import *
 36import sqllog
 37from apiary.tools import timestamp
 38
 39def now(future=0.0):
 40    return timestamp.TimeStamp(time.time() + future)
 41
 42def wait_until(ts):
 43    tn = now()
 44    if ts > tn:
 45        delta = float(ts - tn)
 46        # why was this ever here? --Lex
 47        #if delta > 65.0:
 48        #    raise Exception("trying to wait longer than 65s. now = %s, until = %s"
 49        #        % (str(tn), str(ts)))
 50        time.sleep(delta)
 51
 52mysql_host = 'localhost'
 53mysql_port = 3306
 54mysql_user = 'guest'
 55mysql_passwd = ''
 56mysql_db = 'test'
 57
 58class MySQLWorkerBee(apiary.WorkerBee):
 59    def __init__(self, options, arguments):
 60        apiary.WorkerBee.__init__(self, options, arguments)
 61        self._connect_options = {}
 62        self._connect_options['host'] = options.mysql_host
 63        self._connect_options['port'] = options.mysql_port
 64        self._connect_options['user'] = options.mysql_user
 65        self._connect_options['passwd'] = options.mysql_passwd
 66        self._connect_options['db'] = options.mysql_db
 67        self._connection = None
 68        self._no_mysql = options.no_mysql
 69    
 70    @traced_method
 71    def start(self):
 72        self._error = False
 73        self._errormsg = ''
 74        if self._no_mysql:
 75            return
 76
 77        try:
 78            if self._connection is None and not self._no_mysql:
 79                connection = MySQLdb.connect(**self._connect_options)
 80                self._connection = connection
 81            pass
 82        except Exception, e: # more restrictive error catching?
 83            self._error = True
 84            self._errormsg = "500 " + str(e)
 85        
 86    @traced_method
 87    def event(self, data):
 88        if self._error:
 89            return
 90        try:
 91            (tstr, sql) = data.split("\t", 1)
 92            if not self._asap:
 93                self.log("waiting")
 94                wait_until(timestamp.TimeStamp(tstr))
 95            sql = sql.strip()
 96            if sql:
 97                self.log("executing SQL: " + sql)
 98                self.execute_sql(sql)  
 99        except Exception, e: # more restrictive error catching?
100            self._error = True
101            self._errormsg = "501 " + str(e)
102        
103    @traced_method
104    def end(self):
105        if self._no_mysql:
106            return "200 OK"
107
108        try:
109            if True:
110                self._connection.close()
111                self._connection = None
112        except Exception, e:
113            if not self._error:
114                self._error = True
115                self._errormsg = "502 " + str(e)
116        if self._error:
117            return self._errormsg
118        return '200 OK'
119
120    @traced_method
121    def execute_sql(self, sql):
122        """Execute an SQL statement.
123        
124        Subclasses may override this to alter the SQL before being executed.
125        If so, they should call this inherited version to actually execute
126        the statement. It is acceptable for a sub-class to call this version
127        multiple times.
128        
129        Exceptions are caught in the outer calling function (event())
130        
131        """
132        if self._no_mysql:
133            return
134        cursor = self._connection.cursor()
135        cursor.execute(sql)
136        try:
137            cursor.fetchall()
138        except:
139            pass # not all SQL has data to fetch
140        cursor.close()
141        
142
143class MySQLQueenBee(apiary.QueenBee):
144    def __init__(self, options, arguments):
145        apiary.QueenBee.__init__(self, options, arguments)
146        if options.pickled:
147            self._events = sqllog.input_pickled_events(arguments)
148        else:
149            self._events = sqllog.input_events(arguments)
150        # this builds a sequence of events from the log streams in the
151        # arguments, which come here from the command line
152        self._connections = {}
153        self._tally = {}
154        self._time_scale = 1.0 / options.speedup
155        self._event_start = None
156        self._replay_start = None
157        self._tally_time = time.time() + 15.0
158        self._delay_start = timestamp.TimeStamp(options.prefill)
159        
160    def tally(self, msg):
161        # aggregate these error codes since we see a lot of them (1062/1064)
162        if "Duplicate entry" in msg:
163            msg = '501 (1062, "Duplicate entry for key")'
164        if "You have an error in your SQL syntax" in msg:
165            msg = '501 (1064, "You have an error in your SQL syntax")'
166        self._tally[msg] = self._tally.get(msg, 0) + 1
167        if time.time() > self._tally_time:
168            self.print_tally()
169
170    
171    def print_tally(self):
172        keys = self._tally.keys()
173        keys.sort()
174        print
175        print "       count - message"
176        print "------------   -------------------------------------------"
177        for k in keys:
178            print ("%12d - %s" % (self._tally[k], k))
179        self._tally_time = time.time() + 15.0
180
181        
182    @traced_method
183    def next(self):
184        try:
185            while True:
186                e = self._events.next() # e is a sqllog.Event object
187                if e.state != sqllog.Event.Response:
188                    break
189        except StopIteration:
190            return False
191        
192        if self._event_start is None:
193            self._event_start = e.time
194        t = (e.time - self._event_start) * self._time_scale + self._replay_start + self._delay_start
195        
196        id = e.id
197        if e.state == sqllog.Event.End:
198            if id in self._connections:
199                self.tally("102 End connection")
200                del self._connections[id]
201                self.end(id)
202            else:
203                self.tally("103 Duplicate end")
204        else:
205            if id not in self._connections:
206                self.tally("100 Start connection")
207                s = self._connections[id] = True
208                self.start(id)
209            self.tally("101 Event")
210            self.event(id, str(t) + "\t" + e.body)
211
212        return True
213    
214    def result(self, seq, d):
215        self.tally(d)
216    
217        
218    @traced_method
219    def main(self):
220        t = - time.time()
221        c = - time.clock()
222        self._replay_start = now(1.0)
223        apiary.QueenBee.main(self)
224        c += time.clock()
225        t += time.time()
226
227        print ("Timing: %f process clock, %f wall clock" % (c, t))
228        self.print_tally()
229
230
231# Plugin interface:
232queenbee_cls = MySQLQueenBee
233workerbee_cls = MySQLWorkerBee
234
235
236def add_options(parser):
237    parser.add_option('--no-mysql', default=False, dest='no_mysql', action='store_true',
238                        help="Don't make mysql connections.  Return '200 OK' instead.")
239    parser.add_option('--speedup', default=1.0, dest='speedup', type='float',
240                        help="Time multiple used when replaying query logs.  2.0 means "
241                             "that queries run twice as fast (and the entire run takes "
242                             "half the time the capture ran for).")
243    parser.add_option('--pickled',
244                      default=False, action='store_true',
245                      help='sequence file contains pickled events')
246    parser.add_option('--prefill',
247                      default=0, type="int", metavar="SECONDS",
248                      help="""Prefill the queue with jobs for SECONDS.  Use 
249                          this if the queenbee can't generate jobs fast enough.""")
250    parser.add_option('--mysql-host',
251                        default=mysql_host, metavar='HOST',
252                        help='MySQL server to connect to (default: %default)')
253    parser.add_option('--mysql-port',
254                        default=mysql_port, metavar='PORT',
255                        help='MySQL port to connect on (default: %default)')
256    parser.add_option('--mysql-user',
257                        default=mysql_user, metavar='USER',
258                        help='MySQL user to connect as (default: %default)')
259    parser.add_option('--mysql-passwd',
260                        default=mysql_passwd, metavar='PW',
261                        help='MySQL password to connect with (default: %default)')
262    parser.add_option('--mysql-db',
263                        default=mysql_db, metavar='DB',
264                        help='MySQL database to connect to (default: %default)')
265