From 54b6ad599be8313c6193eb0f9f1e5e1c105c6e00 Mon Sep 17 00:00:00 2001
From: Gabor Boros <gabor.brs@gmail.com>
Date: Tue, 21 Jun 2022 19:02:59 +0200
Subject: [PATCH] feat: reql command format

Co-authored-by: lsabi <13497689+lsabi@users.noreply.github.com>
Signed-off-by: Gabor Boros <gabor.brs@gmail.com>
---
 .coveragerc                       |  2 +-
 Makefile                          |  6 +++---
 pytest.ini                        |  1 +
 rethinkdb/ast.py                  | 10 ++++++++++
 rethinkdb/query.py                | 11 +++++++++++
 tests/integration/test_queries.py | 25 +++++++++++++++++++++++++
 6 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/.coveragerc b/.coveragerc
index 572b9e7..0e5be22 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -7,6 +7,6 @@ omit = *tests*
 
 [report]
 sort = cover
-fail_under = 72
+fail_under = 70
 exclude_lines = pragma: no cover
                 if __name__ == .__main__.:
diff --git a/Makefile b/Makefile
index 266ef45..88f681d 100644
--- a/Makefile
+++ b/Makefile
@@ -17,7 +17,7 @@
 PACKAGE_NAME = rethinkdb
 
 PROTO_FILE_NAME = ql2.proto
-PROTO_FILE_URL = https://raw.githubusercontent.com/rethinkdb/rethinkdb/next/src/rdb_protocol/${PROTO_FILE_NAME}
+PROTO_FILE_URL = https://raw.githubusercontent.com/rethinkdb/rethinkdb/70654faefe29bb0b4f6010fa82bd30a207d014d8/src/rdb_protocol/${PROTO_FILE_NAME}
 TARGET_PROTO_FILE = ${PROTO_FILE_NAME}
 
 FILE_CONVERTER_NAME = ./scripts/convert_protofile.py
@@ -118,12 +118,12 @@ test-unit: ## run unit tests and generate coverage
 
 .PHONY: test-integration
 test-integration: ## run unit tests and generate coverage
-	coverage run -m pytest -m "integration" -vv
+	coverage run -m pytest -m "integration" -m "not v2_5" -vv
 	coverage report
 
 .PHONY: test
 test: ## run all tests and generate coverage
-	coverage run -m pytest -vv
+	coverage run -m pytest -m "not v2_5" -vv
 	coverage report
 	coverage xml
 
diff --git a/pytest.ini b/pytest.ini
index 5c244be..328cdea 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -5,3 +5,4 @@ python_functions=test_*
 markers =
     unit: Select only unit tests
     integration: Select only integration tests
+    v2_5: RethinkDB 2.5+ compatible tests
diff --git a/rethinkdb/ast.py b/rethinkdb/ast.py
index 5b232fd..45e4324 100644
--- a/rethinkdb/ast.py
+++ b/rethinkdb/ast.py
@@ -545,6 +545,9 @@ def to_json_string(self, *args):
     def match(self, *args):
         return Match(self, *args)
 
+    def format(self, *args, **kwargs):
+        return Format(self, *args, **kwargs)
+
     def split(self, *args):
         return Split(self, *args)
 
@@ -1624,6 +1627,13 @@ def __init__(self, *args, **kwargs):
         self.statement = "match"
 
 
+class Format(ReqlMethodQuery):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.term_type = P_TERM.FORMAT
+        self.statement = "format"
+
+
 class ToJsonString(ReqlMethodQuery):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
diff --git a/rethinkdb/query.py b/rethinkdb/query.py
index 918d309..dd2ee46 100644
--- a/rethinkdb/query.py
+++ b/rethinkdb/query.py
@@ -50,6 +50,7 @@
     "error",
     "february",
     "floor",
+    "format",
     "friday",
     "ge",
     "geojson",
@@ -787,6 +788,16 @@ def circle(*arguments, **kwargs):
     return ast.Circle(*arguments, **kwargs)
 
 
+def format(*arguments, **kwargs):  # pylint: disable=redefined-builtin
+    """
+    Format command takes a string as a template and formatting parameters as an
+    object. The parameters in the template string must exist as keys in the
+    object, otherwise, an error raised. The template must be a string literal
+    and cannot be the result of other commands.
+    """
+    return ast.Format(*arguments, **kwargs)
+
+
 row = ast.ImplicitVar()
 
 # Time enum values
diff --git a/tests/integration/test_queries.py b/tests/integration/test_queries.py
index 77717a6..67a0d89 100644
--- a/tests/integration/test_queries.py
+++ b/tests/integration/test_queries.py
@@ -389,6 +389,31 @@ def test_floor(conn):
     assert_test_table(query.floor, conn, scenarios)
 
 
+@pytest.mark.integration
+@pytest.mark.v2_5
+def test_format(conn):
+    scenarios = [
+        Scenario(
+            name="text", args=["hello {name}", {"name": "bob"}], expected="hello bob"
+        ),
+        Scenario(
+            name="numbers", args=["1..2..{count}", {"count": 3}], expected="1..2..3"
+        ),
+        Scenario(
+            name="object",
+            args=["object: {obj}", {"obj": {"foo": "bar"}}],
+            expected='object: {"foo":"bar"}',
+        ),
+        Scenario(
+            name="array",
+            args=["array: {arr}", {"arr": [1, 2, 3]}],
+            expected="array: [1,2,3]",
+        ),
+    ]
+
+    assert_test_table(query.format, conn, scenarios)
+
+
 @pytest.mark.integration
 def test_ge(conn):
     scenarios = [