Update idempotent insert

Author: Yibo-Chen13

example/idempotent/idempotent.py

@@ -0,0 +1,85 @@
+from proton_driver import connect, Client
+from datetime import date
+from time import sleep
+
+
+# Create a test stream
+def create_test_stream(operator, table_name, table_columns):
+    operator.execute(f'DROP STREAM IF EXISTS {table_name};')
+    operator.execute(f'CREATE STREAM {table_name} ({table_columns})')
+
+
+# Use dbapi to implement idempotent insertion
+def use_dbapi():
+    with connect('proton://localhost') as conn:
+        with conn.cursor() as cur:
+            create_test_stream(
+                cur,
+                'test_user',
+                'id int32, name string, birthday date'
+            )
+
+            # Set idempotent_id.
+            cur.set_settings(dict(idempotent_id='batch1'))
+
+            # Insert data into test_user multiple times with the same idempotent_id. # noqa
+            # The query result should contain only the first inserted data.
+            data = [
+                (123456, 'timeplus', date(2024, 10, 24)),
+                (789012, 'stream  ', date(2023, 10, 24)),
+                (135790, 'proton  ', date(2024, 10, 24)),
+                (246801, 'database', date(2024, 10, 24)),
+            ]
+
+            # Execute multiple insert operations.
+            for _ in range(10):
+                cur.execute(
+                    'INSERT INTO test_user (id, name, birthday) VALUES',
+                    data
+                )
+                cur.fetchall()
+
+            # wait for 3 sec to make sure data available in historical store.
+            sleep(3)
+
+            cur.execute('SELECT count() FROM table(test_user)')
+            res = cur.fetchall()
+
+            # Data is inserted only once,so res == (4,).
+            print(res)
+
+
+# Use Client to implement idempotent insertion
+def use_client():
+    cli = Client('localhost', 8463)
+    create_test_stream(cli, 'test_stream', '`i` int,  `v` string')
+
+    setting = {
+        'idempotent_id': 'batch1'
+    }
+
+    data = [
+        (1, 'a'), (2, 'b'), (3, 'c'), (4, 'd'),
+        (5, 'e'), (6, 'f'), (7, 'g'), (8, 'h')
+    ]
+
+    # Execute multiple insert operations.
+    for _ in range(10):
+        cli.execute(
+            'INSERT INTO test_stream (i, v) VALUES',
+            data,
+            settings=setting
+        )
+
+    # wait for 3 sec to make sure data available in historical store.
+    sleep(3)
+
+    res = cli.execute('SELECT count() FROM table(test_stream)')
+
+    # Data is inserted only once,so res == (8,).
+    print(res)
+
+
+if __name__ == "__main__":
+    use_dbapi()  # (4,)
+    use_client()  # (8,)

proton_driver/settings/available.py

@@ -402,4 +402,7 @@
     'format_regexp_escaping_rule': SettingString,
     'format_regexp_skip_unmatched': SettingBool,
     'output_format_enable_streaming': SettingBool,
+    
+    'idempotent_id': SettingString,
+    'enable_idempotent_processing': SettingBool,
 }

tests/test_dbapi.py

@@ -3,7 +3,7 @@
 from contextlib import contextmanager
 import socket
 from unittest.mock import patch
-
+from time import sleep
 from proton_driver import connect
 from proton_driver.dbapi import (
     ProgrammingError, InterfaceError, OperationalError
@@ -159,6 +159,27 @@ def test_execute_insert(self):
             cursor.execute('INSERT INTO test VALUES', [[4]])
             self.assertEqual(cursor.rowcount, 1)
 
+    def test_idempotent_insert(self):
+        with self.created_cursor() as cursor:
+            cursor.execute('CREATE STREAM test (i int, v string)')
+            data = [
+                (123, 'abc'), (456, 'def'), (789, 'ghi'),
+                (987, 'ihg'), (654, 'fed'), (321, 'cba'),
+            ]
+            cursor.set_settings(dict(idempotent_id='batch1'))
+            for _ in range(10):
+                cursor.execute(
+                    'INSERT INTO test (i, v) VALUES',
+                    data
+                )
+                self.assertEqual(cursor.rowcount, 6)
+            sleep(3)
+            rv = cursor.execute('SELECT count(*) FROM table(test)')
+            rv = cursor.fetchall()
+            self.assertEqual(rv, [(6,)])
+
+            cursor.execute('DROP STREAM test')
+
     def test_description(self):
         with self.created_cursor() as cursor:
             self.assertIsNone(cursor.description)

tests/test_insert.py

@@ -1,5 +1,5 @@
 from datetime import date
-
+from time import sleep
 from tests.testcase import BaseTestCase
 from proton_driver import errors
 from proton_driver.errors import ServerException
@@ -148,6 +148,31 @@ def test_insert_return(self):
             )
             self.assertEqual(rv, 5)
 
+    def test_idempotent_insert(self):
+        self.client.execute('CREATE STREAM test (i int, v string)')
+
+        data = [
+            (123, 'abc'), (456, 'def'), (789, 'ghi'),
+            (987, 'ihg'), (654, 'fed'), (321, 'cba'),
+        ]
+
+        setting = {
+            'idempotent_id': 'batch1'
+        }
+
+        for _ in range(20):
+            rv = self.client.execute(
+                'INSERT INTO test (i, v) VALUES',
+                data,
+                settings=setting
+            )
+            self.assertEqual(rv, 6)
+        sleep(3)
+        rv = self.client.execute('SELECT count(*) FROM table(test)')
+        self.assertEqual(rv, [(6, )])
+
+        self.client.execute('DROP STREAM test')
+
 
 class InsertColumnarTestCase(BaseTestCase):
     def test_insert_tuple_ok(self):