ResultSet.one() fails if the row_factory is using a generator

Author: aboudreault

CHANGELOG.rst

@@ -6,6 +6,7 @@ Bug Fixes
 * Improve and fix socket error-catching code in nonblocking-socket reactors (PYTHON-1024)
 * Non-ASCII characters in schema break CQL string generation (PYTHON-1008)
 * Fix OSS driver's virtual table support against DSE 6.0.X and future server releases (PYTHON-1020)
+* ResultSet.one() fails if the row_factory is using a generator (PYTHON-1026)
 
 Other
 -----

cassandra/cluster.py

@@ -4244,7 +4244,14 @@ def one(self):
         you know a query returns a single row. Consider using an iterator if the
         ResultSet contains more than one row.
         """
-        return self._current_rows[0] if self._current_rows else None
+        row = None
+        if self._current_rows:
+            try:
+                row = self._current_rows[0]
+            except TypeError:  # generator object is not subscriptable, PYTHON-1026
+                row = next(iter(self._current_rows))
+
+        return row
 
     def __iter__(self):
         if self._list_mode:

tests/integration/standard/test_row_factories.py

@@ -181,6 +181,32 @@ def test_ordered_dict_factory(self):
         self.assertEqual(result[1]['k'], result[1]['v'])
         self.assertEqual(result[1]['k'], 2)
 
+    def test_generator_row_factory(self):
+        """
+        Test that ResultSet.one() works with a row_factory that contains a generator.
+
+        @since 3.16
+        @jira_ticket PYTHON-1026
+        @expected_result one() returns the first row
+
+        @test_category queries
+        """
+        def generator_row_factory(column_names, rows):
+            return _gen_row_factory(rows)
+
+        def _gen_row_factory(rows):
+            for r in rows:
+                yield r
+
+        session = self.session
+        session.row_factory = generator_row_factory
+
+        session.execute(self.insert1)
+        result = session.execute(self.select)
+        self.assertIsInstance(result, ResultSet)
+        first_row = result.one()
+        self.assertEqual(first_row[0], first_row[1])
+
 
 class NamedTupleFactoryAndNumericColNamesTests(unittest.TestCase):
     """