# Copyright (c) 2014 Wieland Hoffmann
# License: MIT, see LICENSE for details
from __future__ import absolute_import
import amqp
import logging
import pysolr
import urllib2
from . import config
from .schema import SCHEMA
from contextlib import contextmanager
from functools import partial
from json import loads
from sqlalchemy import create_engine
from sqlalchemy.engine.url import URL
from sqlalchemy.orm import sessionmaker
logger = logging.getLogger("sir")
[docs]class SIR_EXIT(Exception):
pass
[docs]class VersionMismatchException(Exception):
def __init__(self, core, expected, actual):
self.core = core
self.expected = expected
self.actual = actual
def __str__(self):
return "%s: Expected %1.1f, got %1.1f" % (self.core,
self.expected,
self.actual)
[docs]def db_session():
"""
Creates a new :class:`sqla:sqlalchemy.orm.session.sessionmaker`.
:rtype: :class:`sqla:sqlalchemy.orm.session.sessionmaker`
"""
cget = partial(config.CFG.get, "database")
cdict = {"username": cget("user")}
for key in ["password", "host", "port"]:
cdict[key] = cget(key)
cdict["database"] = cget("dbname")
e = create_engine(URL("postgresql", **cdict), server_side_cursors=False)
S = sessionmaker(bind=e)
return S
[docs]@contextmanager
def db_session_ctx(Session):
"""
A context manager yielding a database session.
:param sqlalchemy.orm.session.sessionmaker Session:
"""
session = Session()
try:
yield session
session.commit()
except:
session.rollback()
raise
finally:
session.close()
[docs]def solr_connection(core):
"""
Creates a :class:`solr:solr.Solr` connection for the core ``core``.
:param str core:
:raises urllib2.URLError: if a ping to the cores ping handler doesn't
succeed
:rtype: :class:`solr:solr.Solr`
"""
solr_uri = config.CFG.get("solr", "uri")
core_uri = solr_uri + "/" + core
ping_uri = core_uri + "/admin/ping"
logger.debug("Setting up a connection to %s", solr_uri)
logger.debug("Pinging %s", ping_uri)
urllib2.urlopen(ping_uri)
logger.debug("Connection to the Solr core at %s", core_uri)
return pysolr.Solr(core_uri)
[docs]def solr_version_check(core):
"""
Checks that the version of the Solr core ``core`` matches the one in the
schema.
:param str core:
:raises urllib2.URLError: If the Solr core can't be reached
:raises sir.util.VersionMismatchException: If the version in Solr is
different from the supported one
"""
expected_version = SCHEMA[core].version
solr_uri = config.CFG.get("solr", "uri")
u = urllib2.urlopen("%s/%s/schema/version" % (solr_uri, core))
content = loads(u.read())
seen_version = content["version"]
if not seen_version == expected_version:
raise VersionMismatchException(core, expected_version, seen_version)
logger.debug("%s: version %1.1f matches %1.1f", core, expected_version,
seen_version)
[docs]def check_solr_cores_version(cores):
"""
Checks multiple Solr cores for version compatibility
:param [str] cores: The names of the cores
:raises sir.util.VersionMismatchException: If the version in Solr is
different from the supported one
"""
map(solr_version_check, cores)
[docs]def create_amqp_connection():
# type: () -> amqp.Connection
"""
Creates a connection to an AMQP server.
:rtype: :class:`amqp:amqp.connection.Connection`
"""
cget = partial(config.CFG.get, "rabbitmq")
conn = amqp.Connection(
host=cget("host"),
userid=cget("user"),
password=cget("password"),
virtual_host=cget("vhost"),
)
conn.connect()
return conn