From cf5989406aab1d2dc75192ab192942ce0687ded0 Mon Sep 17 00:00:00 2001
From: Endre Karlson <endre.karlson@gmail.com>
Date: Fri, 8 Mar 2013 14:00:10 +0000
Subject: [PATCH] Move DB init code into bs.sqlalchemy

---
 billingstack/sqlalchemy/api.py                | 20 +++++++++++++++++--
 .../storage/impl_sqlalchemy/__init__.py       | 14 ++++---------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/billingstack/sqlalchemy/api.py b/billingstack/sqlalchemy/api.py
index e67147b..f5de14a 100644
--- a/billingstack/sqlalchemy/api.py
+++ b/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)
diff --git a/billingstack/storage/impl_sqlalchemy/__init__.py b/billingstack/storage/impl_sqlalchemy/__init__.py
index 67505b1..cc00167 100644
--- a/billingstack/storage/impl_sqlalchemy/__init__.py
+++ b/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):
         """