source: main/waeup.cas/trunk/waeup/cas/db.py @ 10604

# 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=True)
    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(ticket='%s', user='%s', service='%s', "
            "sso=%s, timestamp=%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