101 lines
2.9 KiB
Python
101 lines
2.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) 2017 Tomek Wójcik <tomek@bthlabs.pl>
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in
|
|
# all copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
|
# THE SOFTWARE.
|
|
#
|
|
|
|
"""
|
|
q3stats.lib.scripts.utils
|
|
=========================
|
|
|
|
This module contains utility functions used by scripts.
|
|
"""
|
|
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
import sys
|
|
|
|
import q3stats.database
|
|
|
|
RET_OK = 0
|
|
RET_ERROR = 99
|
|
|
|
|
|
def pprint(value, stream=None):
|
|
"""JSON-based pretty print. *stream* defaults to ``sys.stdout``."""
|
|
if stream is None:
|
|
stream = sys.stdout
|
|
|
|
stream.write(json.dumps(value, indent=4, sort_keys=True) + "\n")
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def db_session(config):
|
|
"""Context manager for getting the DB session configured for the current
|
|
environment.
|
|
|
|
On exception, the session will be rolled back. The session is always
|
|
properly disposed of, regardless of any errors.
|
|
|
|
**Example**:
|
|
|
|
.. sourcecode:: python
|
|
|
|
def some_function():
|
|
with db_session() as session:
|
|
# Do something with the session.
|
|
session.commit()
|
|
"""
|
|
if q3stats.database.db_engine is None:
|
|
q3stats.database.db_engine = q3stats.database.create_engine(config)
|
|
q3stats.database.db_session = q3stats.database.create_session(
|
|
q3stats.database.db_engine, config
|
|
)
|
|
|
|
session = q3stats.database.db_session
|
|
try:
|
|
yield session
|
|
except:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.remove()
|
|
|
|
|
|
def get_logger(name, args):
|
|
"""Sets up a ``logging.Logger`` instance for script named *name*."""
|
|
logger = logging.getLogger(name)
|
|
|
|
handler = logging.StreamHandler()
|
|
|
|
formatter = logging.Formatter(
|
|
'%(asctime)s ' + name + ': %(levelname)s: %(message)s'
|
|
)
|
|
|
|
handler.setFormatter(formatter)
|
|
logger.addHandler(handler)
|
|
|
|
if args.verbose is True:
|
|
logger.setLevel(logging.DEBUG)
|
|
else:
|
|
logger.setLevel(logging.WARNING)
|
|
|
|
return logger
|