~ubuntu-branches/ubuntu/utopic/gozerbot/utopic

« back to all changes in this revision

Viewing changes to gozerbot/database/alchemy.py

  • Committer: Bazaar Package Importer
  • Author(s): Jeremy Malcolm
  • Date: 2009-09-14 09:00:29 UTC
  • mfrom: (1.1.4 upstream) (3.1.5 sid)
  • Revision ID: james.westby@ubuntu.com-20090914090029-uval0ekt72kmklxw
Tags: 0.9.1.3-3
Changed dependency on python-setuptools to python-pkg-resources
(Closes: #546435) 

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# gozerbot/databse/alchemy.py
 
2
#
 
3
#
 
4
 
 
5
""" alchemy interface. """
 
6
 
 
7
__copyright__ = 'this file is in the public domain'
 
8
 
 
9
## IMPORT SECTION
 
10
 
 
11
# gozerbot imports
 
12
from gozerbot.stats import stats
 
13
from gozerbot.datadir import datadir
 
14
from gozerbot.config import config
 
15
from gozerbot.utils.locking import lockdec
 
16
from gozerbot.utils.log import rlog
 
17
from gozerbot.utils.exception import handle_exception
 
18
 
 
19
# sqlalchemy imports
 
20
from sqlalchemy.ext.declarative import declarative_base
 
21
from sqlalchemy.ext.associationproxy import association_proxy
 
22
from sqlalchemy.ext.orderinglist import ordering_list
 
23
from sqlalchemy import Text, Integer, Sequence, ForeignKey, DateTime
 
24
from sqlalchemy import create_engine, Column, String, Table
 
25
from sqlalchemy.orm import  scoped_session, sessionmaker, relation, eagerload
 
26
from sqlalchemy.orm import create_session as cs
 
27
 
 
28
# basic imports
 
29
import sqlalchemy, thread, os, time, logging
 
30
 
 
31
## END IMPORT
 
32
 
 
33
# debug settings
 
34
if config['debug']:
 
35
    logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
 
36
 
 
37
# locks
 
38
alchemylock = thread.allocate_lock()
 
39
dblocked = lockdec(alchemylock)
 
40
sessionlock = thread.allocate_lock()
 
41
sessionlocked = lockdec(sessionlock)
 
42
querylock = thread.allocate_lock()
 
43
querylocked = lockdec(querylock)
 
44
createlock = thread.allocate_lock()
 
45
createlocked = lockdec(createlock)
 
46
 
 
47
created = []
 
48
 
 
49
@createlocked
 
50
def create_all(plugname='all', base=None):
 
51
    rlog(10, 'alchemy', 'running create_all (%s)' % plugname)
 
52
    if plugname not in created:
 
53
        created.append(plugname)
 
54
        if not base:
 
55
            base = Base
 
56
        base.metadata.create_all()
 
57
    else:
 
58
        rlog(10, 'alchemy', '%s tables already created' % plugname)
 
59
 
 
60
def geturi(ddir=None, mainconfig=None):
 
61
 
 
62
    """  determine database URI from config file """
 
63
 
 
64
    d = ddir or datadir
 
65
 
 
66
    # set config file
 
67
    if mainconfig:
 
68
        config = mainconfig 
 
69
    else:
 
70
        from gozerbot.config import config
 
71
 
 
72
    # if dburi not provided in config file construct it
 
73
    if not config['dburi']:
 
74
 
 
75
        if not 'sqlite' in config['dbtype'] and not 'mysql' in config['dbtype']:
 
76
            dburi = "%s://%s:%s@%s/%s" % (config['dbtype'], config['dbuser'], \
 
77
config['dbpasswd'], config['dbhost'], config['dbname'])
 
78
        elif 'mysql' in config['dbtype']:
 
79
            dburi = "%s://%s:%s@%s/%s?charset=utf8&use_unicode=0" % (config['dbtype'], config['dbuser'], \
 
80
config['dbpasswd'], config['dbhost'], config['dbname'])
 
81
        else:
 
82
            if not os.path.isdir(d + os.sep + 'db'):
 
83
                os.mkdir(d + os.sep + 'db')
 
84
            dburi = "sqlite:///%s/%s" % (ddir or datadir, config['dbname'])
 
85
 
 
86
    else:
 
87
        # dburi found in config
 
88
        dburi = config['dburi']
 
89
 
 
90
        # determine dbtype
 
91
        try:
 
92
            dbtype = dburi.split(':')[0]
 
93
        except:
 
94
            rlog(10, 'alchemy', "can't extract db data from dburi")
 
95
            dbtype = 'unknown'
 
96
 
 
97
        # save dbtype
 
98
        if config['dbtype'] != dbtype:
 
99
            config['dbtype'] = dbtype
 
100
            config.save()
 
101
 
 
102
    return dburi
 
103
 
 
104
def dbstart(ddir=None, mainconfig=None, base=None):
 
105
 
 
106
    """ start the database connection setting Session and engine. """
 
107
 
 
108
    dburi = geturi(ddir, mainconfig)
 
109
 
 
110
    # only show dburi if it doesn't contain a password
 
111
    if '///' in dburi:
 
112
        rlog(10, 'alchemy', 'starting database %s' % dburi)
 
113
    else:
 
114
        rlog(10, 'alchemy', 'starting database')
 
115
 
 
116
    # create engine
 
117
    if 'mysql' in config['dbtype']:
 
118
        engine = create_engine(dburi, strategy='threadlocal', pool_recycle=3600, max_overflow=-1)
 
119
    else:
 
120
        engine = create_engine(dburi, strategy='threadlocal')
 
121
 
 
122
    # setup metadata and session
 
123
    if not base:
 
124
        base = Base
 
125
    base.metadata.bind = engine
 
126
    create_all()
 
127
    rlog(10, 'alchemy', 'done')
 
128
    Session = scoped_session(sessionmaker(autoflush=True))
 
129
    Session.configure(bind=engine)
 
130
    stats.up('alchemy', 'engines')
 
131
 
 
132
    return (Session, engine)
 
133
 
 
134
# vars
 
135
Base = declarative_base()
 
136
#Session, engine = dbstart(datadir)
 
137
Session = engine = None
 
138
 
 
139
def startmaindb(ddir=None, mainconfig=None):
 
140
 
 
141
    """ start the main database. """
 
142
 
 
143
    global Session
 
144
    global engine
 
145
 
 
146
    Session, engine = dbstart(ddir, mainconfig)
 
147
    #pass
 
148
 
 
149
 
 
150
### MODEL
 
151
 
 
152
user = Table('user', Base.metadata,
 
153
    Column('name', String(255), primary_key=True)
 
154
)
 
155
 
 
156
email = Table('email', Base.metadata,
 
157
    Column('name', String(255), ForeignKey(user.c.name), nullable=False),
 
158
    Column('email', String(255), nullable=False),
 
159
    Column('order', Integer, nullable=False)
 
160
)
 
161
 
 
162
class User(Base):
 
163
    __table__ = user
 
164
    _userhosts = relation("UserHost", backref="user", cascade="all, delete-orphan")
 
165
    _perms = relation("Perms", backref="user", cascade="all, delete-orphan")
 
166
    _permits = relation("Permits", backref="user",cascade="all, delete-orphan" )
 
167
    _statuses = relation("Statuses", backref="user", cascade="all, delete-orphan")
 
168
    _pasword = relation("Passwords", backref="user", cascade="all, delete-orphan")
 
169
    _email = relation("Email", backref="user", collection_class=ordering_list('order'),
 
170
                        cascade="all, delete-orphan", order_by=[email.c.order])
 
171
    email = association_proxy('_email', 'email')
 
172
    userhosts = association_proxy('_userhosts', 'userhost')
 
173
    perms = association_proxy('_perms', 'perm')
 
174
    permits = association_proxy('_permits', 'permit')
 
175
    statuses = association_proxy('_statuses', 'status')
 
176
    password = association_proxy('_password', 'passwd')
 
177
 
 
178
class Email(Base):
 
179
    __table__ = email
 
180
    __mapper_args__ = {'primary_key':[email.c.name,email.c.email]}
 
181
 
 
182
    def __init__(self, email):
 
183
        self.email = email
 
184
 
 
185
class UserHost(Base):
 
186
    __tablename__ = 'userhosts'
 
187
    userhost = Column('userhost', String(255), primary_key=True)
 
188
    name = Column('name', String(255), ForeignKey('user.name'), nullable=False)
 
189
 
 
190
    def __init__(self, userhost):
 
191
        self.userhost = userhost
 
192
 
 
193
class Perms(Base):
 
194
    __tablename__ = 'perms'
 
195
    name = Column('name', String(255), ForeignKey('user.name'), nullable=False)
 
196
    perm = Column('perm', String(255), nullable=False)
 
197
    __mapper_args__ = {'primary_key':[name,perm]}
 
198
 
 
199
    def __init__(self, perm):
 
200
        self.perm = perm
 
201
 
 
202
class Permits(Base):
 
203
    __tablename__ = 'permits'
 
204
    name = Column('name', String(255), ForeignKey('user.name'), nullable=False)
 
205
    permit = Column('permit', String(255), nullable=False)
 
206
    __mapper_args__ = {'primary_key':[name,permit]}
 
207
 
 
208
    def __init__(self, permit):
 
209
        self.permit = permit
 
210
 
 
211
class Statuses(Base):    
 
212
    __tablename__ = 'statuses'
 
213
    name = Column('name', String(255), ForeignKey('user.name'), nullable=False)
 
214
    status = Column('status', String(255), nullable=False)
 
215
    __mapper_args__ = {'primary_key':[name,status]}
 
216
 
 
217
    def __init__(self, status):
 
218
        self.status = status
 
219
 
 
220
class Passwords(Base):    
 
221
    __tablename__ = 'passwords'
 
222
    name = Column('name', String(255), ForeignKey('user.name'), primary_key=True)
 
223
    passwd = Column('passwd', String(255), nullable=False)
 
224
 
 
225
    def __init__(self, passwd):
 
226
        self.passwd = passwd
 
227
 
 
228
### END MODEL
 
229
 
 
230
def trans(func, ismethod=True):
 
231
 
 
232
    """ transaction function attribute. """
 
233
 
 
234
    @dblocked
 
235
    def transaction(*args, **kwargs):
 
236
 
 
237
        """ the tranasction wrapper .. works on methods. """
 
238
 
 
239
        arglist = list(args)
 
240
        res = None
 
241
        try:
 
242
            stats.up('alchemy', 'transactions')
 
243
            session = create_session()
 
244
 
 
245
            try:
 
246
                session.begin(subtransactions=True)
 
247
            except sqlalchemy.exc.InvalidRequestError, ex:
 
248
                rlog(10, 'alchemy', 'error %s: %s' % (str(arglist), str(ex)))
 
249
                #Session.remove()
 
250
                close(session)
 
251
                session = create_session()
 
252
                session.begin(subtransactions=True)
 
253
 
 
254
            if not ismethod:
 
255
                arglist.insert(0, session)
 
256
            else:
 
257
                arglist.insert(1, session)
 
258
 
 
259
            res = func(*arglist, **kwargs)
 
260
 
 
261
            session.flush()
 
262
            engine.commit()
 
263
                
 
264
            return res
 
265
 
 
266
        except sqlalchemy.exc.InvalidRequestError, ex:
 
267
            rlog(10, 'alchemy', 'error %s: %s' % (str(arglist), str(ex)))
 
268
            close(session)
 
269
 
 
270
        except sqlalchemy.exc.TimeoutError:
 
271
            rlog(10, 'alchemy', 'timeout occured')
 
272
            close(session)
 
273
 
 
274
        except Exception ,ex:
 
275
            close(session)
 
276
            raise
 
277
 
 
278
        res = func(*arglist, **kwargs)
 
279
 
 
280
        close(session)
 
281
        return res
 
282
 
 
283
    return transaction
 
284
 
 
285
def transfunc(func):
 
286
 
 
287
    """ transaction wrapper for functions. """
 
288
 
 
289
    return trans(func, ismethod=False)
 
290
 
 
291
def rollback(session):
 
292
 
 
293
    """ rollback on provided session. """
 
294
 
 
295
    rlog(10, 'alchemy', 'rollback on %s' % str(session))
 
296
    session.rollback()
 
297
 
 
298
def close(session):
 
299
 
 
300
    """ close provided session. """
 
301
 
 
302
    #rollback(session)
 
303
    session.close()
 
304
    rlog(10, 'alchemy', '%s closed' % str(session))
 
305
 
 
306
#@sessionlocked
 
307
def create_session():
 
308
 
 
309
    """ create a session ready for use. """
 
310
 
 
311
    stats.up('alchemy', 'sessions')
 
312
    session = Session()
 
313
    return session
 
314
 
 
315
@querylocked
 
316
def query(q, session=None):
 
317
 
 
318
    """ do a query on the database. """
 
319
 
 
320
    res = None
 
321
    #if not session:
 
322
    session = create_session()
 
323
    try:
 
324
        stats.up('alchemy', 'query')
 
325
 
 
326
        #try:
 
327
        #    session.begin(subtransactions=True)
 
328
        #except sqlalchemy.exc.InvalidRequestError, ex:
 
329
        #    rlog(10, 'alchemy', 'error %s: %s' % (str(q), str(ex)))
 
330
            #Session.remove()
 
331
        #    close(session)
 
332
        #    session = create_session()
 
333
        #    session.begin(subtransactions=True)
 
334
 
 
335
        res = session.query(q)
 
336
        return res
 
337
 
 
338
    except sqlalchemy.exc.TimeoutError:
 
339
        rlog(10, 'alchemy', 'timeout occured')
 
340
        close(session)
 
341
 
 
342
    except Exception ,ex:
 
343
        close(session)
 
344
        raise
 
345
 
 
346
    session = create_session()
 
347
    res = session.query(q)
 
348
    return res
 
349
 
 
350
def getuser(userhost, session=None):
 
351
 
 
352
    """ get a user based on userhost. """
 
353
 
 
354
    stats.up('alchemy', 'getuser')
 
355
    session = create_session()
 
356
    try:
 
357
        session.begin(subtransactions=True)
 
358
    except sqlalchemy.exc.InvalidRequestError, ex:
 
359
        Session.remove()
 
360
        #close(session)
 
361
        session = create_session()
 
362
        session.begin(subtransactions=True)
 
363
      
 
364
    try:
 
365
        user = query(UserHost, session).filter_by(userhost=userhost).first()
 
366
        if user:
 
367
            session = create_session()
 
368
            res = query(User, session).filter_by(name=user.name.lower()).first()
 
369
            if res: 
 
370
                return res
 
371
    except KeyError:
 
372
        rlog(10, 'alchemy', 'error getting %s user: %s' % (str(userhost), str(ex)))
 
373
        rollback(session)
 
374
        raise
 
375
    except sqlalchemy.orm.exc.ConcurrentModificationError, ex:
 
376
        rlog(10, 'alchemy', 'error getting %s user: %s' % (str(userhost), str(ex)))
 
377
        rollback(session)
 
378
        raise
 
379
    except sqlalchemy.exc.InvalidRequestError, ex:
 
380
        rlog(10, 'alchemy', 'error getting %s user: %s' % (str(userhost), str(ex)))
 
381
        rollback(session)
 
382
        raise
 
383
 
 
384
def byname(name , session=None):
 
385
 
 
386
    """ get a users based on name. """
 
387
    if not session:
 
388
        session = create_session()
 
389
    try:
 
390
        res = query(User, session).filter_by(name=name.lower()).first()
 
391
        return res
 
392
    except KeyError:
 
393
        rlog(10, 'alchemy', 'error getting %s user: %s' % (str(name), str(ex)))
 
394
        rollback(session)
 
395
    except sqlalchemy.exc.InvalidRequestError, ex:
 
396
        rlog(10, 'alchemy', 'error getting %s user: %s' % (str(name), str(ex)))
 
397
        rollback(session)
 
398
 
 
399
@transfunc
 
400
def dbupgrade(session, mainconfig=None):
 
401
 
 
402
    """ upgrade the database. """
 
403
 
 
404
    #time.sleep(10)
 
405
    print 'upgrading users'
 
406
    users = session.query(UserHost).all()
 
407
    upgraded = []
 
408
 
 
409
    # populate the User table
 
410
    for user in users:
 
411
        name = user.name
 
412
        if name in upgraded:
 
413
            continue
 
414
        try:
 
415
            if not byname(name):
 
416
                newuser = User(name=name)
 
417
                session.add(newuser)
 
418
            upgraded.append(name)
 
419
        except sqlalchemy.exc.IntegrityError, ex:
 
420
            pass
 
421
        except:
 
422
            handle_exception()
 
423
 
 
424
    session.commit()
 
425
    print "upgraded: %s" % ' .. '.join(upgraded) 
 
426
    print 'upgrading email table'
 
427
    from gozerbot.database.db import Db
 
428
 
 
429
    # upgrade email table
 
430
    try:
 
431
        db = Db(config=mainconfig)
 
432
        if db.dbtype == 'mysql':
 
433
            db.execute("ALTER TABLE email ADD COLUMN email.order INT")
 
434
        else:
 
435
            db.execute("ALTER TABLE email ADD COLUMN 'order' INT")
 
436
    except Exception, ex:
 
437
        if 'already exists' in str(ex) or 'duplicate column name' in \
 
438
str(ex).lower():
 
439
            pass
 
440
        else:
 
441
            handle_exception()