def setup(cls, store=None, dataset=None, level=None, logger=None, name=None,
fmt=None, reset=False, size=10 * 1024 * 1024, defaults=None, exit_hook=False):
"""
Args:
dataset (str): the name of the dataset
level (int): set any logging.INFO, logging.ERROR, logging.DEBUG, defaults to INFO
size (int): maxium size in bytes (defaults to 1MB)
target (Omega): omega instance to create the dataset in, defaults to the default om
size (int): the maximum size of the log in bytes (capped), defaults to 1MB, set to -1
for
name (str): the logger name, defaults to omegaml, set to __name__ for current module
reset (bool): recreate the logging dataset
fmt (str): the format specification
exit_hook (bool): when True attach the logger to the system exception handler
"""
import omegaml as om
import logging
effective_level = logger.getEffectiveLevel() if logger else logging.INFO
logger_name = name or 'omegaml'
level = level or effective_level
store = store or om.setup().datasets
defaults = defaults or store.defaults
dataset = dataset or defaults.OMEGA_LOG_DATASET
fmt = fmt or defaults.OMEGA_LOG_FORMAT
# setup handler and logger
logger = logger or logging.getLogger(logger_name)
logger.setLevel(level)
collection = _setup_logging_dataset(store, dataset, logger=logger, size=size, reset=reset)
formatter = logging.Formatter(fmt)
handler = OmegaLoggingHandler(store, dataset, collection, level=level)
handler.setFormatter(formatter)
logger.addHandler(handler)
if exit_hook:
_attach_sysexcept_hook(logger)
return handler
def load_sql(om=None, kind=SQLAlchemyBackend.KIND):
"""
load ipython sql magic, loading all sql alchemy connections
Usage:
!pip install ipython-sql
# prepare some connection, insert some data
df = pd.DataFrame(...)
om.datasets.put('sqlite:///test.db', 'testdb')
om.datasets.put(df, 'testdb', table='foobar', insert=True)
from omegaml.backends.sqlalchemy import load_sql
load_sql()
# list registered connections
%sql
omsql://testdb
# run queries
%sql omsql://testdb select * from foobar
See Also:
https://pypi.org/project/ipython-sql/
Args:
om (Omega): optional, specify omega instance, defaults to om.setup()
kind (str): the backend's kind, used to find connections to register
with the sql magic
Returns:
None
"""
from unittest.mock import MagicMock
from IPython import get_ipython
import omegaml as om
from sql.connection import Connection
class ConnectionShim:
# this is required to trick sql magic into accepting existing connection objects
# (by default sql magic expects connection strings, not connection objects)
def __init__(self, url, conn):
self.session = conn
self.metadata = MagicMock()
self.metadata.bind.url = url
self.dialect = getattr(conn, 'dialect', 'omsql')
# load sql magic
ipython = get_ipython()
ipython.magic('load_ext sql')
# load registered sqlalchemy datasets
om = om or om.setup()
for ds in om.datasets.list(kind=kind, raw=True):
cnxstr = 'omsql://{ds.name}'.format(**locals())
conn = om.datasets.get(ds.name, raw=True)
Connection.connections[cnxstr] = ConnectionShim(cnxstr, conn)
ipython.magic('sql')
def before_all(context):
# setup omegaml
import omegaml as om
use_fixture(splinter_browser, context)
# set url and admin password
context.web_url = os.environ.get('OMEGA_URL', 'http://localhost:5000')
context.jynb_url = os.environ.get('JUPYTER_URL', 'http://localhost:8888')
# setup environment
context.debug = os.environ.get('BEHAVE_DEBUG', False)
defaults = settings()
context.om = om.setup()
context.nbfiles = os.environ.get('BEHAVE_NBFILES', './docs/source/nb')
def get_omega(args):
from omegaml import setup, _base_config
from omegaml.client.cloud import setup_from_config
config_file = args.get('--config')
local_runtime = args.get('--local-runtime')
if local_runtime:
_base_config.OMEGA_LOCAL_RUNTIME = True
bucket = args.get('--bucket')
if config_file:
om = setup_from_config(config_file)
else:
om = setup()
return om[bucket]
def test_config_from_apikey(self):
"""
Test an Omega instance can be created from user specific configs
"""
import omegaml as om
from omegaml.util import settings
# check we get default without patching
defaults = settings()
setup = om.setup
with patch.object(defaults, 'OMEGA_MONGO_URL') as mock:
defaults.OMEGA_MONGO_URL = 'foo'
om.setup()
self.assertEqual(om.datasets.mongo_url, 'foo')
# reset om.datasets to restored defaults
setup()
self.assertNotEqual(om.datasets.mongo_url, 'foo')
# now test we can change the default through config
# we patch the actual api call to avoid having to set up the user db
# the objective here is to test get_omega_from_apikey
with patch('omegaml.client.userconf.get_user_config_from_api') as mock:
mock.return_value = {
'objects': [
{
'data': {
'OMEGA_MONGO_URL': 'updated-foo'
}
}
]
}
with patch.object(defaults, 'OMEGA_MONGO_URL') as mock:
from omegaml.client.userconf import get_omega_from_apikey
defaults.OMEGA_MONGO_URL = 'foo'
om.setup()
self.assertEqual(om.datasets.mongo_url, 'foo')
om = get_omega_from_apikey('foo', 'bar')
self.assertEqual(om.datasets.mongo_url, 'updated-foo')
setup()
self.assertNotEqual(om.datasets.mongo_url, 'foo')
def __init__(self, store=None, dataset=None, collection=None, level='INFO',
size=1 * 1024 * 1024, defaults=None, name='simple'):
import omegaml as om
self.store = store or om.setup().datasets
self.defaults = defaults or store.defaults
self.dsname = dataset or self.defaults.OMEGA_LOG_DATASET
self._dataset = None
self._level = None
self.setLevel(level)
self._is_setup = False
self._collection = collection
self.size = size
self._name = name
def get_omega(args, require_config=False):
from omegaml import setup, _base_config
from omegaml.client.cloud import setup_from_config
config_file = args.get('--config')
# deprecated, use --local
local_runtime = args.get('--local-runtime')
local = args.get('--local')
if local or local_runtime:
_base_config.OMEGA_LOCAL_RUNTIME = True
bucket = args.get('--bucket')
if config_file or require_config:
try:
om = setup_from_config(config_file)
except Exception as e:
msg = (f'Config file could not be found due to {e}. Specify as --config or set '
'OMEGA_CONFIG_FILE env variable')
raise ValueError(msg)
else:
om = setup()
if local:
om.runtime.mode(local=True)
return om[bucket] if bucket else om # for speed
def setup(cls,
store=None,
dataset=None,
level=None,
logger=None,
fmt=None,
reset=False,
size=10 * 1024 * 1024,
defaults=None):
"""
Args:
dataset (str): the name of the dataset
level (int): set any logging.INFO, logging.ERROR, logging.DEBUG
size (int): maxium size in bytes (defaults to 1MB)
target (Omega): omega instance to create the dataset in, defaults to the default om
size (int): the maximum size of the log in bytes (capped), defaults to 1MB, set to -1
for
reset (bool): recreate the logging dataset
fmt (str): the format specification
"""
import omegaml as om
import logging
level = level or logging.INFO
store = store or om.setup().datasets
defaults = defaults or store.defaults
dataset = dataset or defaults.OMEGA_LOG_DATASET
fmt = fmt or defaults.OMEGA_LOG_FORMAT
# setup handler and logger
logger = logger or logging.getLogger('omegaml')
collection = _setup_logging_dataset(store,
dataset,
logger=logger,
size=size,
reset=reset)
formatter = logging.Formatter(fmt)
handler = OmegaLoggingHandler(store, dataset, collection, level=level)
handler.setFormatter(formatter)
logger.addHandler(handler)
return handler
def test_config_from_apikey(self):
"""
Test an Omega instance can be created from user specific configs
"""
import omegaml as om
from omegaml.util import settings
# check we get default without patching
from omegaml import _base_config as _real_base_config
with patch('omegaml._base_config', new=BareObj) as defaults:
# link callbacks used by get_omega_from_api_key
_real_base_config.update_from_obj(_real_base_config,
attrs=defaults)
defaults.update_from_dict = _real_base_config.update_from_dict
defaults.load_user_extensions = lambda *args, **kwargs: None
defaults.load_framework_support = lambda *args, **kwargs: None
setup = om.setup
defaults.MY_OWN_SETTING = 'foo'
settings(reload=True)
om = om.setup()
self.assertEqual(om.defaults.MY_OWN_SETTING, 'foo')
# reset om.datasets to restored defaults
om = setup()
self.assertNotEqual(om.datasets.mongo_url, 'foo')
# now test we can change the default through config
# we patch the actual api call to avoid having to set up the user db
# the objective here is to test get_omega_from_apikey
with patch('omegaml.client.userconf.get_user_config_from_api') as mock:
mock.return_value = {
'objects': [{
'data': {
'OMEGA_MONGO_URL': 'updated-foo',
'OMEGA_MY_OWN_SETTING': 'updated-foo',
'OMEGA_CELERY_CONFIG': {
'TEST_SETTING': 'pickle',
}
}
}]
}
with patch('omegaml._base_config', new=BareObj) as defaults:
from omegaml.client.userconf import get_omega_from_apikey
# link callbacks used by get_omega_from_api_key
_real_base_config.update_from_obj(_real_base_config,
attrs=defaults)
defaults.update_from_dict = _real_base_config.update_from_dict
defaults.load_user_extensions = lambda *args, **kwargs: None
defaults.load_framework_support = lambda *args, **kwargs: None
defaults.OMEGA_MY_OWN_SETTING = 'foo'
om = get_omega_from_apikey('foo', 'bar')
self.assertEqual(om.defaults.OMEGA_MY_OWN_SETTING,
'updated-foo')
self.assertEqual(om.datasets.mongo_url, 'updated-foo')
self.assertEqual(
om.defaults.OMEGA_CELERY_CONFIG['TEST_SETTING'], 'pickle')
# test that all default values are still there, i.e. the OMEGA_CELERY_CONFIG was updated, not replaced
for real_k, real_v in _real_base_config.OMEGA_CELERY_CONFIG.items(
):
self.assertIn(real_k, om.defaults.OMEGA_CELERY_CONFIG)
# restore defaults
defaults = settings(reload=True)
om = setup()
self.assertNotEqual(om.datasets.mongo_url, 'foo')
from __future__ import absolute_import
'''
make sure Celery is correctly configured
see http://chriskief.com/2013/11/15/celery-3-1-with-django-django-celery-rabbitmq-and-macports/
'''
import os
from celery import Celery
from omegaml import settings as omsettings
# get rid of celery's Django compatibility mode
os.environ['DJANGO_SETTINGS_MODULE'] = ''
try:
import omegaml as om
# see if we can get an instance from the configuration env
om = om.setup()
defaults = om.defaults
except:
# otherwise just load local defaults
defaults = omsettings()
app = Celery('omegaml')
app.config_from_object(defaults.OMEGA_CELERY_CONFIG)
app.autodiscover_tasks(defaults.OMEGA_CELERY_IMPORTS, related_name='tasks')
def get_omega_from_apikey(cls, auth=None):
# return the omega instance for the given task authentication
from omegaml import setup
om = setup()
return om
def _omega(self):
if self._omega_instance is None:
bucket = flask.request.headers.get('bucket')
self._omega_instance = om.setup()[bucket]
return self._omega_instance
def connection(ctx):
import omegaml as om
ctx.om = om.setup()
print(ctx.om)
assert ctx.om is not None
def connection(ctx):
import omegaml as om
ctx.feature.om = om.setup()
assert ctx.feature.om is not None
def _omega(self):
bucket = flask.request.headers.get('bucket')
return om.setup()[bucket]
