Implement strict checking in one; #13

Author: chadwhitacre

postgres.py

@@ -23,12 +23,11 @@
     >>> db.run("INSERT INTO foo VALUES ('baz')")
     >>> db.run("INSERT INTO foo VALUES ('buz')")
 
-Use :py:meth:`~postgres.Postgres.one` to fetch one result:
+Use :py:meth:`~postgres.Postgres.one` to fetch exactly one result:
 
-    >>> db.one("SELECT * FROM foo ORDER BY bar")
+    >>> db.one("SELECT * FROM foo WHERE bar='baz'")
     {'bar': 'baz'}
-    >>> db.one("SELECT * FROM foo WHERE bar='blam'")
-    >>> # None
+
 
 Use :py:meth:`~postgres.Postgres.rows` to fetch all results:
 
@@ -177,6 +176,19 @@ def url_to_dsn(url):
     return dsn
 
 
+# Exceptions
+# ==========
+
+class NotOne(Exception):
+    def __init__(self, rowcount):
+        self.rowcount = rowcount
+    def __str__(self):
+        return "Got {} rows instead of 1.".format(self.rowcount)
+
+class TooFew(NotOne): pass
+class TooMany(NotOne): pass
+
+
 # The Main Event
 # ==============
 
@@ -186,6 +198,8 @@ class Postgres(object):
     :param unicode url: A ``postgres://`` URL or a `PostgreSQL connection string <http://www.postgresql.org/docs/current/static/libpq-connect.html>`_
     :param int minconn: The minimum size of the connection pool
     :param int maxconn: The minimum size of the connection pool
+    :param strict_one: The default :py:attr:`strict` parameter for :py:meth:`~postgres.Postgres.one`
+    :type strict_one: :py:class:`bool`
 
     This is the main object that :py:mod:`postgres` provides, and you should
     have one instance per process for each PostgreSQL database your process
@@ -220,15 +234,20 @@ class Postgres(object):
 
     """
 
-    def __init__(self, url, minconn=1, maxconn=10):
+    def __init__(self, url, minconn=1, maxconn=10, strict_one=None):
         if url.startswith("postgres://"):
             dsn = url_to_dsn(url)
+
         self.pool = ConnectionPool( minconn=minconn
                                   , maxconn=maxconn
                                   , dsn=dsn
                                   , connection_factory=Connection
                                    )
 
+        if strict_one not in (True, False, None):
+            raise ValueError("strict_one must be True, False, or None.")
+        self.strict_one = strict_one
+
     def run(self, sql, parameters=None):
         """Execute a query and discard any results.
 
@@ -245,21 +264,48 @@ def run(self, sql, parameters=None):
         with self.get_cursor() as cursor:
             cursor.execute(sql, parameters)
 
-    def one(self, sql, parameters=None):
+    def one(self, sql, parameters=None, strict=None):
         """Execute a query and return a single result.
 
         :param unicode sql: the SQL statement to execute
         :param parameters: the bind parameters for the SQL statement
         :type parameters: dict or tuple
+        :param strict: whether to raise when there isn't exactly one result
+        :type strict: :py:class:`bool`
         :returns: :py:class:`dict` or :py:const:`None`
+        :raises: :py:exc:`~postgres.TooFew` or :py:exc:`~postgres.TooMany`
+
+        By default, :py:attr:`strict` ends up evaluating to :py:class:`True`,
+        in which case we raise :py:exc:`~postgres.TooFew` or
+        :py:exc:`~postgres.TooMany` if the number of rows returned isn't
+        exactly one. You can override this behavior per-call with the
+        :py:attr:`strict` argument here, or globally by passing
+        :py:attr:`strict_one` to the :py:class:`~postgres.Postgres`
+        constructor. If you use both, the :py:attr:`strict` argument here wins.
 
         >>> row = db.one("SELECT * FROM foo WHERE bar='baz'"):
         >>> print(row["bar"])
         baz
 
         """
+        if strict not in (True, False, None):
+            raise ValueError("strict must be True, False, or None.")
+
+        if strict is None:
+            if self.strict_one is None:
+                strict = True               # library default
+            else:
+                strict = self.strict_one    # user default
+
         with self.get_cursor() as cursor:
             cursor.execute(sql, parameters)
+
+            if strict:
+                if cursor.rowcount < 1:
+                    raise TooFew(cursor.rowcount)
+                elif cursor.rowcount > 1:
+                    raise TooMany(cursor.rowcount)
+
             return cursor.fetchone()
 
     def rows(self, sql, parameters=None):

tests.py

@@ -3,12 +3,15 @@
 import os
 from unittest import TestCase
 
-from postgres import Postgres
+from postgres import Postgres, TooFew, TooMany
 
 
 DATABASE_URL = os.environ['DATABASE_URL']
 
 
+# harnesses
+# =========
+
 class WithSchema(TestCase):
 
     def setUp(self):
@@ -18,6 +21,7 @@ def setUp(self):
 
     def tearDown(self):
         self.db.run("DROP SCHEMA IF EXISTS public CASCADE")
+        del self.db
 
 
 class WithData(WithSchema):
@@ -29,6 +33,9 @@ def setUp(self):
         self.db.run("INSERT INTO foo VALUES ('buz')")
 
 
+# db.run
+# ======
+
 class TestRun(WithSchema):
 
     def test_run_runs(self):
@@ -44,16 +51,126 @@ def test_run_inserts(self):
         assert actual == 1
 
 
-class TestOneAndRows(WithData):
+# db.one
+# ======
+# With all the combinations of strict_one and strict, we end up with a number
+# of tests here. Since the behavior of the one method with a strict parameter
+# of True or False is expected to be the same regardless of what strict_one is
+# set to, we can write those once and then use the TestOne TestCase as the base
+# class for other TestCases that vary the strict_one attribute. The TestOne
+# tests will be re-run in each new context.
 
-    def test_one_fetches_the_first_one(self):
-        actual = self.db.one("SELECT * FROM foo ORDER BY bar")
+class TestNotOneException(WithData):
+
+    def test_TooFew_message_is_helpful(self):
+        try:
+            self.db.one("SELECT * FROM foo WHERE bar='blah'", strict=True)
+        except TooFew, exc:
+            actual = str(exc)
+            assert actual == "Got 0 rows instead of 1."
+
+    def test_TooMany_message_is_helpful(self):
+        try:
+            self.db.one("SELECT * FROM foo", strict=True)
+        except TooMany, exc:
+            actual = str(exc)
+            assert actual == "Got 2 rows instead of 1."
+
+
+class TestOne(WithData):
+
+    def test_with_strict_True_one_raises_TooFew(self):
+        self.assertRaises( TooFew
+                         , self.db.one
+                         , "SELECT * FROM foo WHERE bar='blah'"
+                         , strict=True
+                          )
+
+    def test_with_strict_True_one_fetches_the_one(self):
+        actual = self.db.one("SELECT * FROM foo WHERE bar='baz'", strict=True)
+        assert actual == {"bar": "baz"}
+
+    def test_with_strict_True_one_raises_TooMany(self):
+        self.assertRaises( TooMany
+                         , self.db.one
+                         , "SELECT * FROM foo"
+                         , strict=True
+                          )
+
+
+    def test_with_strict_False_one_returns_None_if_theres_none(self):
+        actual = self.db.one("SELECT * FROM foo WHERE bar='nun'", strict=False)
+        assert actual is None
+
+    def test_with_strict_False_one_fetches_the_first_one(self):
+        actual = self.db.one("SELECT * FROM foo ORDER BY bar", strict=False)
+        assert actual == {"bar": "baz"}
+
+
+class TestOne_StrictOneNone(TestOne):
+
+    def setUp(self):
+        WithData.setUp(self)
+        self.db.strict_one = None
+
+    def test_one_raises_TooFew(self):
+        self.assertRaises( TooFew
+                         , self.db.one
+                         , "SELECT * FROM foo WHERE bar='nun'"
+                          )
+
+    def test_one_returns_one(self):
+        actual = self.db.one("SELECT * FROM foo WHERE bar='baz'")
         assert actual == {"bar": "baz"}
 
-    def test_one_returns_None_if_theres_none(self):
-        actual = self.db.one("SELECT * FROM foo WHERE bar='blam'")
+    def test_one_raises_TooMany(self):
+        self.assertRaises(TooMany, self.db.one, "SELECT * FROM foo")
+
+
+class TestOne_StrictOneFalse(TestOne):
+
+    def setUp(self):
+        WithData.setUp(self)
+        self.db.strict_one = False
+
+    def test_one_returns_None(self):
+        actual = self.db.one("SELECT * FROM foo WHERE bar='nun'")
         assert actual is None
 
+    def test_one_returns_one(self):
+        actual = self.db.one("SELECT * FROM foo WHERE bar='baz'")
+        assert actual == {"bar": "baz"}
+
+    def test_one_returns_first_one(self):
+        actual = self.db.one("SELECT * FROM foo ORDER BY bar")
+        assert actual == {"bar": "baz"}
+
+
+class TestOne_StrictOneTrue(TestOne):
+
+    def setUp(self):
+        WithData.setUp(self)
+        self.db.strict_one = True
+
+    def test_one_raises_TooFew(self):
+        self.assertRaises( TooFew
+                         , self.db.one
+                         , "SELECT * FROM foo WHERE bar='nun'"
+                          )
+
+    def test_one_returns_one(self):
+        actual = self.db.one("SELECT * FROM foo WHERE bar='baz'")
+        assert actual == {"bar": "baz"}
+
+    def test_one_raises_TooMany(self):
+        self.assertRaises(TooMany, self.db.one, "SELECT * FROM foo")
+
+
+# db.rows
+# =======
+
+class TestRows(WithData):
+
     def test_rows_fetches_all_rows(self):
         actual = self.db.rows("SELECT * FROM foo ORDER BY bar")
         assert actual == [{"bar": "baz"}, {"bar": "buz"}]
@@ -68,6 +185,9 @@ def test_bind_parameters_as_tuple_work(self):
         assert actual == [{"bar": "baz"}]
 
 
+# db.get_cursor
+# =============
+
 class TestCursor(WithData):
 
     def test_get_cursor_gets_a_cursor(self):
@@ -89,6 +209,9 @@ def test_we_can_use_cursor_closed(self):
         assert not actual
 
 
+# db.get_transaction
+# ==================
+
 class TestTransaction(WithData):
 
     def test_get_transaction_gets_a_transaction(self):
@@ -125,6 +248,9 @@ class Heck(Exception): pass
         assert actual == [{"bar": "baz"}, {"bar": "buz"}]
 
 
+# db.get_connection
+# =================
+
 class TestConnection(WithData):
 
     def test_get_connection_gets_a_connection(self):