ffsync_ynh/sources/syncserver/staticnode.py

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
"""
Simple node-assignment backend using a single, static node.

This is a greatly-simplified node-assignment backend.  It keeps user records
in an SQL database, but does not attempt to do any node management.  All users
are implicitly assigned to a single, static node.

XXX TODO: move this into the tokenserver repo.

"""
import time
import urlparse
from mozsvc.exceptions import BackendError

from sqlalchemy import Column, Integer, String, BigInteger, Index
from sqlalchemy import create_engine, Table, MetaData
from sqlalchemy.pool import QueuePool
from sqlalchemy.sql import text as sqltext
from sqlalchemy.exc import IntegrityError

from tokenserver.assignment import INodeAssignment
from zope.interface import implements


metadata = MetaData()


users = Table(
    "users",
    metadata,
    Column("uid", Integer(), primary_key=True, autoincrement=True,
           nullable=False),
    Column("service", String(32), nullable=False),
    Column("email", String(255), nullable=False),
    Column("generation", BigInteger(), nullable=False),
    Column("client_state", String(32), nullable=False),
    Column("created_at", BigInteger(), nullable=False),
    Column("replaced_at", BigInteger(), nullable=True),
    Index('lookup_idx', 'email', 'service', 'created_at'),
    Index('clientstate_idx', 'email', 'service', 'client_state', unique=True),
)


_GET_USER_RECORDS = sqltext("""\
select
    uid, generation, client_state
from
    users
where
    email = :email
and
    service = :service
order by
    created_at desc, uid desc
limit
    20
""")


_CREATE_USER_RECORD = sqltext("""\
insert into
    users
    (service, email, generation, client_state, created_at, replaced_at)
values
    (:service, :email, :generation, :client_state, :timestamp, NULL)
""")


_UPDATE_GENERATION_NUMBER = sqltext("""\
update
    users
set
    generation = :generation
where
    service = :service and email = :email and
    generation < :generation and replaced_at is null
""")


_REPLACE_USER_RECORDS = sqltext("""\
update
    users
set
    replaced_at = :timestamp
where
    service = :service and email = :email
    and replaced_at is null and created_at < :timestamp
""")


def get_timestamp():
    """Get current timestamp in milliseconds."""
    return int(time.time() * 1000)


class StaticNodeAssignment(object):
    implements(INodeAssignment)

    def __init__(self, sqluri, node_url, **kw):
        self.sqluri = sqluri
        self.node_url = node_url
        self.driver = urlparse.urlparse(sqluri).scheme.lower()
        sqlkw = {
            "logging_name": "syncserver",
            "connect_args": {},
            "poolclass": QueuePool,
            "pool_reset_on_return": True,
        }
        if self.driver == "sqlite":
            # We must mark it as safe to share sqlite connections between
            # threads.  The pool will ensure there's on race conditions.
            sqlkw["connect_args"]["check_same_thread"] = False
            # If using a :memory: database, we must use a QueuePool of size
            # 1 so that a single connection is shared by all threads.
            if urlparse.urlparse(sqluri).path.lower() in ("/", "/:memory:"):
                sqlkw["pool_size"] = 1
                sqlkw["max_overflow"] = 0
        self._engine = create_engine(sqluri, **sqlkw)
        users.create(self._engine, checkfirst=True)

    def get_user(self, service, email):
        params = {'service': service, 'email': email}
        res = self._engine.execute(_GET_USER_RECORDS, **params)
        try:
            row = res.fetchone()
            if row is None:
                return None
            # The first row is the most up-to-date user record.
            user = {
                'email': email,
                'uid': row.uid,
                'node': self.node_url,
                'generation': row.generation,
                'client_state': row.client_state,
                'old_client_states': {}
            }
            # Any subsequent rows are due to old client-state values.
            row = res.fetchone()
            while row is not None:
                user['old_client_states'][row.client_state] = True
                row = res.fetchone()
            return user
        finally:
            res.close()

    def allocate_user(self, service, email, generation=0, client_state=''):
        params = {
            'service': service, 'email': email, 'generation': generation,
            'client_state': client_state, 'timestamp': get_timestamp()
        }
        try:
            res = self._engine.execute(_CREATE_USER_RECORD, **params)
        except IntegrityError:
            raise
            return self.get_user(service, email)
        else:
            res.close()
            return {
                'email': email,
                'uid': res.lastrowid,
                'node': self.node_url,
                'generation': generation,
                'client_state': client_state,
                'old_client_states': {}
            }

    def update_user(self, service, user, generation=None, client_state=None):
        if client_state is None:
            # uid can stay the same, just update the generation number.
            if generation is not None:
                params = {
                    'service': service,
                    'email': user['email'],
                    'generation': generation,
                }
                res = self._engine.execute(_UPDATE_GENERATION_NUMBER, **params)
                res.close()
                user['generation'] = max(generation, user['generation'])
        else:
            # reject previously-seen client-state strings.
            if client_state == user['client_state']:
                raise BackendError('previously seen client-state string')
            if client_state in user['old_client_states']:
                raise BackendError('previously seen client-state string')
            # need to create a new record for new client_state.
            if generation is not None:
                generation = max(user['generation'], generation)
            else:
                generation = user['generation']
            now = get_timestamp()
            params = {
                'service': service, 'email': user['email'],
                'generation': generation, 'client_state': client_state,
                'timestamp': now,
            }
            try:
                res = self._engine.execute(_CREATE_USER_RECORD, **params)
            except IntegrityError:
                user.update(self.get_user(service, user['email']))
            else:
                self.get_user(service, user['email'])
                user['uid'] = res.lastrowid
                user['generation'] = generation
                user['old_client_states'][user['client_state']] = True
                user['client_state'] = client_state
                res.close()
            # mark old records as having been replaced.
            # if we crash here, they are unmarked and we may fail to
            # garbage collect them for a while, but the active state
            # will be undamaged.
            params = {
                'service': service, 'email': user['email'], 'timestamp': now
            }
            res = self._engine.execute(_REPLACE_USER_RECORDS, **params)
            res.close()