From 5bbac8f17c52ff4df83373d8aea4847d0016ff3c Mon Sep 17 00:00:00 2001
From: Kenneth Reitz <me@kennethreitz.org>
Date: Sat, 6 Feb 2016 23:50:34 -0500
Subject: [PATCH] cleanup

---
 proto.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/proto.py b/proto.py
index 585f8f7..21a2c47 100644
--- a/proto.py
+++ b/proto.py
@@ -70,9 +70,17 @@ class ResultSet(object):
 class Database(object):
     """A Database connection."""
 
-    def __init__(self, conn_str):
-        self._conn_str = conn_str
-        self.db = psycopg2.connect(conn_str, cursor_factory=RealDictCursor)
+    def __init__(self, db_url=None):
+
+        # If no db_url was provided, fallback to $DATABASE_URL.
+        if not db_url and DATABASE_URL:
+            db_url = DATABASE_URL
+
+        if not db_url:
+            raise ValueError('You must provide a db_url.')
+
+        self.db_url = db_url
+        self.db = psycopg2.connect(self.db_url, cursor_factory=RealDictCursor)
 
         # Enable hstore if it's available.
         self._enable_hstore()