From 086c575079e039d000c3dde7073b028601618c5b Mon Sep 17 00:00:00 2001
From: mdumandag <metin@hazelcast.com>
Date: Fri, 17 Sep 2021 12:36:52 +0300
Subject: [PATCH] Make `SqlRow` support `[]` operator

This is mainly an enhancement for easier access for columns in the row.
Now, the `SqlRow` implements the `__getitem__` magic method.

Implementation of the method simply performs a type check and calls
`get_object_with_index` if the type of the passed parameter is an
integer type. For any other type, it calls `get_object`. Hence,
it will throw for any other type than `int` or `str`.

Also, put a piece of information about this new way to the documentation
& code samples.
---
 ...sing_python_client_with_hazelcast_imdg.rst |  7 +++++-
 examples/sql/sql_example.py                   | 14 ++++++++++-
 hazelcast/sql.py                              | 25 ++++++++++++++++++-
 .../backward_compatible/sql_test.py           | 16 ++++++++++++
 4 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/docs/using_python_client_with_hazelcast_imdg.rst b/docs/using_python_client_with_hazelcast_imdg.rst
index d43833d8e7..46ec3ae48f 100644
--- a/docs/using_python_client_with_hazelcast_imdg.rst
+++ b/docs/using_python_client_with_hazelcast_imdg.rst
@@ -1567,6 +1567,8 @@ The following code prints names of the employees whose age is less than 30:
 
     for row in result:
         name = row.get_object("name")
+        # or, you can use the [] operator
+        # name = row["name"]
         print(name)
 
 
@@ -1681,7 +1683,10 @@ it:
 
     with client.sql.execute(select_query) as result:
         for row in result:
-            print(row.get_object("name"))
+            name = row.get_object("name")
+            # or, you can use the [] operator
+            # name = row["name"]
+            print(name)
 
 
 **"SELECT *" Queries**
diff --git a/examples/sql/sql_example.py b/examples/sql/sql_example.py
index 1cd555f4e3..201d3bb55c 100644
--- a/examples/sql/sql_example.py
+++ b/examples/sql/sql_example.py
@@ -67,6 +67,11 @@ def __repr__(self):
         # Get the object with the column index
         is_active = row.get_object_with_index(is_active_index)
 
+        # Rows can also be used similar to lists
+        # with integer column indexes
+
+        # is_active = row[is_active_index]
+
         print(name, age, is_active)
 
 # Construct a statement object to control the properties of the query
@@ -86,13 +91,20 @@ def __repr__(self):
         key = row.get_object("__key")
         age = row.get_object("age")
 
+        # Rows can also be used similar to dictionaries
+        # with string column names
+
+        # key = row["__key"]
+        # age = row["age"]
+
         print(key, age)
 
 # Parameters can be passed directly in the basic execution syntax
 result = client.sql.execute("SELECT this FROM customers WHERE age > ? AND age < ?", 30, 40)
 
 for row in result:
-    customer = row.get_object("this")
+    # Access columns with [] operator
+    customer = row["this"]
     print(customer)
 
 # Query can be closed explicitly
diff --git a/hazelcast/sql.py b/hazelcast/sql.py
index 9a11c12011..8ec1e26aa1 100644
--- a/hazelcast/sql.py
+++ b/hazelcast/sql.py
@@ -521,7 +521,24 @@ def __repr__(self):
 
 
 class SqlRow(object):
-    """One of the rows of an SQL query result."""
+    """One of the rows of an SQL query result.
+
+    The columns of the rows can be retrieved using
+
+    - :func:`get_object` with column name.
+    - :func:`get_object_with_index` with column index.
+
+    Apart from these methods, the row objects can also be treated as a ``dict``
+    or ``list`` and columns can be retrieved using the ``[]`` operator.
+
+    If an integer value is passed to the ``[]`` operator, it will implicitly
+    call the :func:`get_object_with_index` and return the result.
+
+    For any other type passed into the the ``[]`` operator, :func:`get_object`
+    will be called. Note that, :func:`get_object` expects ``str`` values.
+    Hence, the ``[]`` operator will raise error for any type other than integer
+    and string.
+    """
 
     __slots__ = ("_row_metadata", "_row", "_deserialize_fn")
 
@@ -609,6 +626,12 @@ def metadata(self):
         """SqlRowMetadata: The row metadata."""
         return self._row_metadata
 
+    def __getitem__(self, item):
+        if isinstance(item, six.integer_types):
+            return self.get_object_with_index(item)
+
+        return self.get_object(item)
+
     def __repr__(self):
         return "[%s]" % ", ".join(
             "%s %s=%s"
diff --git a/tests/integration/backward_compatible/sql_test.py b/tests/integration/backward_compatible/sql_test.py
index ae5309a666..efef873e3e 100644
--- a/tests/integration/backward_compatible/sql_test.py
+++ b/tests/integration/backward_compatible/sql_test.py
@@ -572,6 +572,22 @@ def test_lazy_deserialization(self):
             with self.assertRaises(HazelcastSqlError):
                 row.get_object("this")
 
+    def test_rows_as_dict_or_list(self):
+        skip_if_client_version_older_than(self, "5.0")
+
+        self._create_mapping("VARCHAR")
+        entry_count = 20
+
+        def value_factory(v):
+            return "value-%s" % v
+
+        self._populate_map(entry_count, value_factory)
+
+        expected = [(i, value_factory(i)) for i in range(entry_count)]
+        with self.client.sql.execute('SELECT __key, this FROM "%s"' % self.map_name) as result:
+            # Verify that both row[integer] and row[string] works
+            six.assertCountEqual(self, expected, [(row[0], row["this"]) for row in result])
+
 
 @unittest.skipIf(
     compare_client_version("4.2") < 0, "Tests the features added in 4.2 version of the client"