Merge pull request #24 from newtdb/connection-fixes

Connection fixes

CHANGES.rst

@@ -1,10 +1,13 @@
 Changes
 =======
 
-0.6.1 (unreleased)
+0.7.0 (unreleased)
 ------------------
 
-- Nothing changed yet.
+- Fixed: newt.db.search.query_data didn't accept an ordinary ZODB connection.
+
+- The methods in newt.db.search that accept a connection now
+  accept a database object (in addition to a Newt or ZODB connection).
 
 
 0.6.0 (2017-04-19)

src/newt/db/search.py

@@ -203,7 +203,7 @@ def create_text_index(conn, fname, D, C=None, B=None, A=None, config=None):
     connection, but a separate connection is used, so it's execution
     is independent of the current transaction.
     """
-    conn, cursor = conn._storage.ex_connect()
+    conn, cursor = _storage(conn).ex_connect()
     sql = create_text_index_sql(fname, D, C, B, A, config)
     try:
         cursor.execute(sql)
@@ -233,7 +233,8 @@ def query_data(conn, query, *args, **kw):
             raise TypeError("Only positional or keyword arguments"
                             " may be provided, not both.")
         args = kw
-    cursor = conn._connection._storage.ex_cursor()
+
+    cursor = read_only_cursor(conn)
     try:
         cursor.execute(query, args)
         result = list(cursor)
@@ -310,6 +311,12 @@ def where_batch(conn, query_tail, args, batch_start, batch_size=None):
                         args, batch_start, batch_size)
 
 
+def _storage(conn):
+    try:
+        return conn._storage
+    except AttributeError:
+        return conn._p_jar._storage
+
 def read_only_cursor(conn):
     """Get a database cursor for reading.
 
@@ -321,4 +328,4 @@ def read_only_cursor(conn):
 
     The caller must close the returned cursor after use.
     """
-    return conn._storage.ex_cursor()
+    return _storage(conn).ex_cursor()

src/newt/db/tests/testsearch.py

@@ -220,6 +220,26 @@ def test_create_text_index_standalone(self):
             set(self.conn.where("txt(state) @@ 'bar | green'")),
             )
 
+    def test_create_text_index_db_object(self):
+        from .. import search
+        conn = self.conn.root()
+        search.create_text_index(conn, 'txt', 'text')
+        self.store('a', text='foo bar')
+        self.store('b', text='foo baz')
+        self.store('c', text='green eggs and spam')
+        self.assertEqual(
+            set((self.conn.root.a, self.conn.root.b)),
+            set(self.conn.where("txt(state) @@ 'foo'")),
+            )
+        self.assertEqual(
+            set((self.conn.root.a, )),
+            set(self.conn.where("txt(state) @@ 'foo & bar'")),
+            )
+        self.assertEqual(
+            set((self.conn.root.a, self.conn.root.c)),
+            set(self.conn.where("txt(state) @@ 'bar | green'")),
+            )
+
     def test_query_data(self):
         from .. import search
         self.store('a', text='foo bar')
@@ -242,6 +262,26 @@ def test_query_data(self):
                  where state @> '{"text": "foo bar"}'""")
              ])
 
+        # Make sure we can search using a ZODB connection:
+        self.assertEqual(
+            [[1]],
+            [list(map(int, r)) for r in
+             search.query_data(
+                 self.conn._connection,
+                 """select zoid from newt
+                 where state @> '{"text": "foo bar"}'""")
+             ])
+
+        # For good mesaue, we can search with a persistent object:
+        self.assertEqual(
+            [[1]],
+            [list(map(int, r)) for r in
+             search.query_data(
+                 self.conn._connection.root(),
+                 """select zoid from newt
+                 where state @> '{"text": "foo bar"}'""")
+             ])
+
 expect_simple_text = """\
 create or replace function mytext(state jsonb) returns tsvector as $$
 declare