# components to store tickets in a database # we use sqlite3 import time from sqlalchemy import ( create_engine, Column, Integer, Float, String, Boolean, MetaData) from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker, scoped_session from sqlalchemy.pool import StaticPool Base = declarative_base() class ServiceTicket(Base): __tablename__ = 'service_tickets' id = Column(Integer, primary_key=True) ticket = Column(String(255), nullable=False) ts = Column(Float, nullable=False) service = Column(String(2048), nullable=True) user = Column(String(512), nullable=False) sso = Column(Boolean, nullable=False, default=True) def __init__(self, ticket, user, service=None, sso=True, timestamp=None): if timestamp is None: timestamp = time.time() self.ticket = ticket self.ts = timestamp self.service = service self.user = user self.sso = sso def __repr__(self): return "ServiceTicket('%s', '%s', '%s', %s, %s)" % ( self.ticket, self.user, self.service, self.sso, self.ts) class LoginTicket(Base): __tablename__ = 'login_tickets' ticket = Column(String(255), primary_key=True) ts = Column(Float, nullable=False) def __init__(self, ticket, timestamp=None): if timestamp is None: timestamp = time.time() self.ticket = ticket self.ts = timestamp def __repr__(self): return "LoginTicket('%s', %s)" % (self.ticket, self.ts) class TicketGrantingCookie(Base): __tablename__ = 'ticket_granting_cookies' value = Column(String(255), primary_key=True) ts = Column(Float, nullable=False) def __init__(self, value, timestamp=None): if timestamp is None: timestamp = time.time() self.value = value self.ts = timestamp def __repr__(self): return "TicketGrantingCookie('%s', %s)" % (self.value, self.ts) Session = scoped_session(sessionmaker()) class DB(object): """Abstract database to make data persistent. Creates a database identified by `connection_string` if it does not exist and initializes sessions. """ @property def session(self): return Session def __init__(self, connection_string): args = {} if connection_string in ('sqlite:///', 'sqlite:///:memory:'): # make sure all threads access the same memory args = dict(connect_args={'check_same_thread': False}, poolclass=StaticPool) self.engine = create_engine(connection_string, **args) self.metadata = MetaData() Base.metadata.create_all(self.engine) Session.configure(bind=self.engine) def add(self, item): """Insert `item` into database. `item` must be an instance of the database content types defined in this module, normally some `Ticket` instance. """ Session.add(item) Session.commit() def delete(self, item): """Delete `item` from database. `item must be an instance of the database content types defined in this module, normally some `Ticket` instance. """ Session.delete(item) Session.commit() def query(self, *args, **kw): return Session.query(*args, **kw) class DBSessionContext(object): """A context manager providing database sessions. Creates a new (threadlocal) database session when entering and tears down the session after leaving the context. Tearing down includes committing transactions and closing the connection. Meant to be used as a wrapper for web request handlers, such, that a session is created when a request comes in and released when the response is ready to be delivered. This context manager does *not* catch any exceptions. """ def __enter__(self): return Session() def __exit__(self, *args, **kw): if args or kw: Session.rollback() Session.remove() return False