Move DB init code into bs.sqlalchemy

billingstack/sqlalchemy/api.py

@@ -2,10 +2,26 @@ from sqlalchemy import or_
 from sqlalchemy.orm import exc
 
 from billingstack import exceptions
-from billingstack.sqlalchemy.model_base import ModelBase
+from billingstack.openstack.common import log
+from billingstack.sqlalchemy import model_base, session
+
+
+LOG = log.getLogger(__name__)
 
 
 class HelpersMixin(object):
+    def setup(self, config_group):
+        self.session = session.get_session(config_group)
+        self.engine = session.get_engine(config_group)
+
+    def setup_schema(self):
+        """ Semi-Private Method to create the database schema """
+        self.base.metadata.create_all(self.session.bind)
+
+    def teardown_schema(self):
+        """ Semi-Private Method to reset the database schema """
+        self.base.metadata.drop_all(self.session.bind)
+
     def _save(self, obj, save=True):
         if not save:
             return obj
@@ -100,7 +116,7 @@ class HelpersMixin(object):
         :param obj: ID or instance / ref of the object
         :param cls: The class to run self._get on if obj is not a ref
         """
-        if isinstance(obj, ModelBase):
+        if isinstance(obj, model_base.ModelBase):
             return obj
         elif isinstance(obj, basestring) and cls:
             return self._get(cls, obj)

billingstack/storage/impl_sqlalchemy/__init__.py

@@ -17,9 +17,9 @@ from billingstack.openstack.common import log as logging
 from billingstack import exceptions
 from billingstack import utils as common_utils
 from billingstack.sqlalchemy import utils as db_utils, api
+from billingstack.sqlalchemy.session import SQLOPTS
 from billingstack.storage import base
 from billingstack.storage.impl_sqlalchemy import models
-from billingstack.storage.impl_sqlalchemy.session import get_session, SQLOPTS
 
 
 LOG = logging.getLogger(__name__)
@@ -42,16 +42,10 @@ class Connection(base.Connection, api.HelpersMixin):
     """
     SQLAlchemy connection
     """
+    base = models.BASE
+
     def __init__(self, config_group):
-        self.session = get_session(config_group)
-
-    def setup_schema(self):
-        """ Semi-Private Method to create the database schema """
-        models.BASE.metadata.create_all(self.session.bind)
-
-    def teardown_schema(self):
-        """ Semi-Private Method to reset the database schema """
-        models.BASE.metadata.drop_all(self.session.bind)
+        self.setup(config_group)
 
     def set_properties(self, obj, properties, cls=None, rel_attr='properties', purge=False):
         """