From 00a410ad1a1444dc8b87185916ebc26aff47b01b Mon Sep 17 00:00:00 2001
From: slene <vslene@gmail.com>
Date: Mon, 10 Mar 2014 20:50:54 +0800
Subject: [PATCH] orm add GetDB api #433

---
 orm/db_alias.go | 20 ++++++++++++++++++--
 orm/orm_test.go |  9 +++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/orm/db_alias.go b/orm/db_alias.go
index 22066514..6a6623cc 100644
--- a/orm/db_alias.go
+++ b/orm/db_alias.go
@@ -168,7 +168,7 @@ func addAliasWthDB(aliasName, driverName string, db *sql.DB) (*alias, error) {
 	}
 
 	if dataBaseCache.add(aliasName, al) == false {
-		return nil, fmt.Errorf("db name `%s` already registered, cannot reuse", aliasName)
+		return nil, fmt.Errorf("DataBase alias name `%s` already registered, cannot reuse", aliasName)
 	}
 
 	return al, nil
@@ -239,7 +239,7 @@ func SetDataBaseTZ(aliasName string, tz *time.Location) error {
 	if al, ok := dataBaseCache.get(aliasName); ok {
 		al.TZ = tz
 	} else {
-		return fmt.Errorf("DataBase name `%s` not registered\n", aliasName)
+		return fmt.Errorf("DataBase alias name `%s` not registered\n", aliasName)
 	}
 	return nil
 }
@@ -260,3 +260,19 @@ func SetMaxOpenConns(aliasName string, maxOpenConns int) {
 		fun.Call([]reflect.Value{reflect.ValueOf(maxOpenConns)})
 	}
 }
+
+// Get *sql.DB from registered database by db alias name.
+// Use "default" as alias name if you not set.
+func GetDB(aliasNames ...string) (*sql.DB, error) {
+	var name string
+	if len(aliasNames) > 0 {
+		name = aliasNames[0]
+	} else {
+		name = "default"
+	}
+	if al, ok := dataBaseCache.get(name); ok {
+		return al.DB, nil
+	} else {
+		return nil, fmt.Errorf("DataBase of alias name `%s` not found\n", name)
+	}
+}
diff --git a/orm/orm_test.go b/orm/orm_test.go
index 90f2ec80..69f2fc86 100644
--- a/orm/orm_test.go
+++ b/orm/orm_test.go
@@ -139,6 +139,15 @@ func throwFailNow(t *testing.T, err error, args ...interface{}) {
 	}
 }
 
+func TestGetDB(t *testing.T) {
+	if db, err := GetDB(); err != nil {
+		throwFailNow(t, err)
+	} else {
+		err = db.Ping()
+		throwFailNow(t, err)
+	}
+}
+
 func TestSyncDb(t *testing.T) {
 	RegisterModel(new(Data), new(DataNull))
 	RegisterModel(new(User))