feat: Sender(auto_flush) with watermark.

Author: amunra

src/questdb/ingress.pyx

@@ -266,14 +266,13 @@ cdef class TimestampNanos:
 
 
 cdef class Sender
+cdef class Buffer
 
 
-cdef int may_flush_on_row_complete(
-    line_sender_buffer* buf,
-    Sender sender) except -1:
+cdef int may_flush_on_row_complete(Buffer buffer, Sender sender) except -1:
     if sender._auto_flush_enabled:
-        if line_sender_buffer_size(buf) >= sender._auto_flush_watermark:
-            print('w00t')
+        if len(buffer) >= sender._auto_flush_watermark:
+            sender.flush(buffer)
 
 
 cdef class Buffer:
@@ -476,9 +475,7 @@ cdef class Buffer:
         if self._row_complete_sender != None:
             sender = PyWeakref_GetObject(self._row_complete_sender)
             if sender != NULL:
-                may_flush_on_row_complete(
-                    self._impl,
-                    <Sender><object>sender)
+                may_flush_on_row_complete(self, <Sender><object>sender)
 
     cdef inline int _at_ts(self, TimestampNanos ts) except -1:
         cdef line_sender_error* err = NULL
@@ -772,7 +769,7 @@ cdef class Sender:
     cdef line_sender* _impl
     cdef Buffer _buffer
     cdef bint _auto_flush_enabled
-    cdef size_t _auto_flush_watermark
+    cdef ssize_t _auto_flush_watermark
     cdef object _init_capacity
     cdef object _max_name_len
 
@@ -785,9 +782,9 @@ cdef class Sender:
             tuple auth=None,
             object tls=False,
             object read_timeout=None,
-            int init_capacity=65536,
+            int init_capacity=65536,  # 64KiB
             int max_name_len=127,
-            object auto_flush=32768):
+            object auto_flush=64512):  # 63KiB
         cdef line_sender_error* err = NULL
 
         cdef line_sender_utf8 host_utf8
@@ -887,6 +884,10 @@ cdef class Sender:
         self._auto_flush_enabled = not not auto_flush
         self._auto_flush_watermark = int(auto_flush) \
             if self._auto_flush_enabled else 0
+        if self._auto_flush_watermark < 0:
+            raise ValueError(
+                'auto_flush_watermark must be >= 0, '
+                f'not {self._auto_flush_watermark}')
 
     def new_buffer(self):
         """

test/test.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 import sys
-from turtle import clear
 sys.dont_write_bytecode = True
 import os
 import unittest
@@ -222,8 +221,38 @@ def test_auto_flush(self):
             with qi.Sender('localhost', server.port, auto_flush=4) as sender:
                 server.accept()
                 sender.row('tbl1', symbols={'sym1': 'val1'})
-                # msgs = server.recv()
-                # self.assertEqual(msgs, [b'tbl1,sym1=val1', b'tbl1,sym2=val2'])
+                self.assertEqual(len(sender), 0)  # auto-flushed buffer.
+                msgs = server.recv()
+                self.assertEqual(msgs, [b'tbl1,sym1=val1'])
+
+    def test_immediate_auto_flush(self):
+        with Server() as server:
+            with qi.Sender('localhost', server.port, auto_flush=True) as sender:
+                server.accept()
+                sender.row('tbl1', symbols={'sym1': 'val1'})
+                self.assertEqual(len(sender), 0)  # auto-flushed buffer.
+                msgs = server.recv()
+                self.assertEqual(msgs, [b'tbl1,sym1=val1'])
+
+    def test_dont_auto_flush(self):
+        with Server() as server:
+            with qi.Sender('localhost', server.port, auto_flush=0) as sender:
+                server.accept()
+                while len(sender) < 131072:  # 128KiB
+                    sender.row('tbl1', symbols={'sym1': 'val1'})
+                msgs = server.recv()
+                self.assertEqual(msgs, [])
+
+    def test_dont_flush_on_exception(self):
+        with Server() as server:
+            with self.assertRaises(RuntimeError):
+                with qi.Sender('localhost', server.port) as sender:
+                    server.accept()
+                    sender.row('tbl1', symbols={'sym1': 'val1'})
+                    self.assertEqual(str(sender), 'tbl1,sym1=val1\n')
+                    raise RuntimeError('Test exception')
+            msgs = server.recv()
+            self.assertEqual(msgs, [])
 
 
 if __name__ == '__main__':