New auto_flush flag logic, no tests yet

Author: amunra

examples/random_data.py

@@ -1,12 +1,12 @@
-from questdb.ingress import Sender, TimestampNanos
+from questdb.ingress import Sender, TimestampNanos, AutoFlush
 import random
 import uuid
 import time
 
 
 def example(host: str = 'localhost', port: int = 9009):
     table_name: str = str(uuid.uuid1())
-    watermark = 1024  # Flush if the internal buffer exceeds 1KiB
+    watermark = AutoFlush.ByteCount(1024)  # Flush if the internal buffer exceeds 1KiB
     with Sender(host=host, port=port, auto_flush=watermark) as sender:
         total_rows = 0
         last_flush = time.monotonic()

src/questdb/dataframe.pxi

@@ -2,12 +2,14 @@
 
 cdef struct auto_flush_t:
     line_sender* sender
+    auto_flush_mode_t mode
     size_t watermark
 
 
 cdef auto_flush_t auto_flush_blank() noexcept nogil:
     cdef auto_flush_t af
     af.sender = NULL
+    af.mode = auto_flush_disabled
     af.watermark = 0
     return af
 
@@ -2094,8 +2096,14 @@ cdef void_int _dataframe_handle_auto_flush(
     cdef line_sender_error* marker_err
     cdef bint flush_ok
     cdef bint marker_ok
-    if (af.sender == NULL) or (line_sender_buffer_size(ls_buf) < af.watermark):
+    if (af.sender == NULL) or (af.mode == auto_flush_disabled):
         return 0
+    elif af.mode == auto_flush_row_count:
+        if line_sender_buffer_row_count(ls_buf) < af.watermark:
+            return 0
+    elif af.mode == auto_flush_byte_count:
+        if line_sender_buffer_size(ls_buf) < af.watermark:
+            return 0
 
     # Always temporarily release GIL during a flush.
     had_gil = _ensure_doesnt_have_gil(gs)

src/questdb/ingress.pyx

@@ -496,7 +496,11 @@ cdef class Buffer
 
 
 cdef void_int may_flush_on_row_complete(Buffer buffer, Sender sender) except -1:
-    if sender._auto_flush_enabled:
+    cdef bint flush = False
+    if sender._auto_flush_mode == auto_flush_row_count:
+        if line_sender_buffer_row_count(buffer._impl) >= sender._auto_flush_watermark:
+            sender.flush(buffer)
+    elif sender._auto_flush_mode == auto_flush_byte_count:
         if len(buffer) >= sender._auto_flush_watermark:
             sender.flush(buffer)
 
@@ -1208,6 +1212,46 @@ _FLUSH_FMT = ('{} - See https://py-questdb-client.readthedocs.io/en/'
     '/troubleshooting.html#inspecting-and-debugging-errors#flush-failed')
 
 
+class AutoFlush:
+    """
+    Auto flush mode.
+
+    Use one of ``AutoFlush.Disabled``, ``AutoFlush.RowCount`` or
+    ``AutoFlush.ByteCount``.
+    """
+
+    class Disabled:
+        """
+        Auto flush disabled.
+        Call ``flush()`` manually.
+        """
+        pass
+
+    class RowCount:
+        """
+        Auto flush after a number of rows.
+        """
+        def __init__(self, value: int):
+            if value < 0:
+                raise ValueError('value must be a positive integer.')
+            self.value = value
+
+    class ByteCount:
+        """
+        Auto flush after a number of bytes.
+        """
+        def __init__(self, value: int):
+            if value < 0:
+                raise ValueError('value must be a positive integer.')
+            self.value = value
+
+
+cdef enum auto_flush_mode_t:
+    auto_flush_disabled,
+    auto_flush_row_count,
+    auto_flush_byte_count,
+
+
 cdef class Sender:
     """
     A sender is a client that inserts rows into QuestDB via the ILP protocol.
@@ -1265,19 +1309,25 @@ cdef class Sender:
     You can control this behavior by setting the ``auto_flush`` argument.
 
     .. code-block:: python
+        from questdb.ingress import AutoFlush
 
-        # Never flushes automatically.
-        sender = Sender('localhost', 9009, auto_flush=False)
-        sender = Sender('localhost', 9009, auto_flush=None) # Ditto.
-        sender = Sender('localhost', 9009, auto_flush=0)  # Ditto.
+        # Flushes automatically after 600 rows (the default).
+        sender = Sender('localhost', 9009)
+        sender = Sender('localhost', 9009, auto_flush=AutoFlush.RowCount(600)) # Ditto.
+
+        # Flushes automatically after each row.
+        sender = Sender('localhost', 9009, auto_flush=AutoFlush.RowCount(1))
 
         # Flushes automatically when the buffer reaches 1KiB.
-        sender = Sender('localhost', 9009, auto_flush=1024)
+        sender = Sender('localhost', 9009, auto_flush=AutoFlush.ByteCount(1024))
 
-        # Flushes automatically after every row.
-        sender = Sender('localhost', 9009, auto_flush=True)
-        sender = Sender('localhost', 9009, auto_flush=1)  # Ditto.
+        # Disabled: Never flushes automatically.
+        sender = Sender('localhost', 9009, auto_flush=AutoFlush.Disabled)
 
+    When auto-flushing is disabled, you must call ``sender.flush()`` manually.
+    Note that when exiting a ``with sender:`` block, the sender will
+    automatically flush the buffer. This is *not* disabled by setting
+    ``auto_flush=AutoFlush.Disabled``.
 
     **Authentication and TLS Encryption**
 
@@ -1348,9 +1398,9 @@ cdef class Sender:
     * ``max_name_length`` (``int``): Maximum length of a table or column name.
       *See Buffer's constructor for more details.*
 
-    * ``auto_flush`` (``bool`` or ``int``): Whether to automatically flush the
-      buffer when it reaches a certain byte-size watermark.
-      *Default: 64512 (63KiB).*
+    * ``auto_flush``: Whether to automatically flush the
+      buffer when it reaches a certain row-count or byte-count watermark.
+      *Default: 600 rows.*
       *See above for details.*
 
     **HTTP-only keyword-only constructor arguments for the Sender(..)**
@@ -1359,9 +1409,9 @@ cdef class Sender:
       the rows in the batch are for the same table.
       Setting ``transactional=True`` will prevent flushing batches of rows
       with mixed table names.
-      To fully control transactions, you also need to set ``auto_flush=False``
-      or buffered lines may be flushed automatically and thus split across
-      multiple transactions.
+      To fully control transactions, you also need to set
+      ``auto_flush=AutoFlush.Disabled`` or buffered lines may be flushed
+      automatically and thus split across multiple transactions.
       If ``transactional=False`` (default), the client will send the batch
       as-is to the server even if it contains rows for multiple tables.
       In such case the server may end up committing some rows and not others.
@@ -1390,8 +1440,8 @@ cdef class Sender:
     cdef line_sender_opts* _opts
     cdef line_sender* _impl
     cdef Buffer _buffer
-    cdef bint _auto_flush_enabled
-    cdef ssize_t _auto_flush_watermark
+    cdef auto_flush_mode_t _auto_flush_mode
+    cdef size_t _auto_flush_watermark
     cdef size_t _init_capacity
     cdef size_t _max_name_len
 
@@ -1407,7 +1457,7 @@ cdef class Sender:
             uint64_t read_timeout=15000,
             uint64_t init_capacity=65536,  # 64KiB
             uint64_t max_name_len=127,
-            object auto_flush=64512,   # 63KiB
+            object auto_flush=None,   # AutoFlush.RowCount(600)
             bint transactional=False,
             uint32_t max_retries=3,
             uint64_t retry_interval=100,  # milliseconds
@@ -1443,6 +1493,8 @@ cdef class Sender:
 
         cdef qdb_pystr_buf* b
 
+        cdef ssize_t auto_flush_watermark = 0
+
         self._opts = NULL
         self._impl = NULL
 
@@ -1532,13 +1584,30 @@ cdef class Sender:
         line_sender_opts_retry_interval(self._opts, retry_interval)
         line_sender_opts_min_throughput(self._opts, min_throughput)
 
-        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:
+
+        # Parse `auto_flush` argument.
+        if auto_flush is None:
+            self._auto_flush_mode = auto_flush_row_count
+            auto_flush_watermark = 600
+        elif auto_flush is AutoFlush.Disabled:
+            self._auto_flush_mode = auto_flush_disabled
+            auto_flush_watermark = 0
+        elif isinstance(auto_flush, AutoFlush.RowCount):
+            self._auto_flush_mode = auto_flush_row_count
+            auto_flush_watermark = auto_flush.value
+        elif isinstance(auto_flush, AutoFlush.ByteCount):
+            self._auto_flush_mode = auto_flush_byte_count
+            auto_flush_watermark = auto_flush.value
+        else:
+            raise TypeError(
+                'auto_flush must be AutoFlush.Disabled, '
+                'AutoFlush.RowCount or AutoFlush.ByteCount, '
+                f'not {_fqn(type(auto_flush))}')
+        if auto_flush_watermark < 0:
             raise ValueError(
-                'auto_flush_watermark must be >= 0, '
-                f'not {self._auto_flush_watermark}')
+                'auto_flush watermark must be >= 0, '
+                f'not {auto_flush_watermark}')
+        self._auto_flush_watermark = auto_flush_watermark
 
         qdb_pystr_buf_clear(b)
 
@@ -1678,8 +1747,9 @@ cdef class Sender:
         may have been transmitted to the server already.
         """
         cdef auto_flush_t af = auto_flush_blank()
-        if self._auto_flush_enabled:
+        if self._auto_flush_mode:
             af.sender = self._impl
+            af.mode = self._auto_flush_mode
             af.watermark = self._auto_flush_watermark
         _dataframe(
             af,

test/test.py

@@ -359,7 +359,7 @@ def test_independent_buffer(self):
 
     def test_auto_flush(self):
         with Server() as server:
-            with qi.Sender('localhost', server.port, auto_flush=4) as sender:
+            with qi.Sender('localhost', server.port, auto_flush=qi.AutoFlush.ByteCount(4)) as sender:
                 server.accept()
                 sender.row('tbl1', symbols={'sym1': 'val1'})
                 self.assertEqual(len(sender), 0)  # auto-flushed buffer.
@@ -368,7 +368,7 @@ def test_auto_flush(self):
 
     def test_immediate_auto_flush(self):
         with Server() as server:
-            with qi.Sender('localhost', server.port, auto_flush=True) as sender:
+            with qi.Sender('localhost', server.port, auto_flush=qi.AutoFlush.RowCount(1)) as sender:
                 server.accept()
                 sender.row('tbl1', symbols={'sym1': 'val1'})
                 self.assertEqual(len(sender), 0)  # auto-flushed buffer.
@@ -377,7 +377,7 @@ def test_immediate_auto_flush(self):
 
     def test_auto_flush_on_closed_socket(self):
         with Server() as server:
-            with qi.Sender('localhost', server.port, auto_flush=True) as sender:
+            with qi.Sender('localhost', server.port, auto_flush=qi.AutoFlush.RowCount(1)) as sender:
                 server.accept()
                 server.close()
                 exp_err = 'Could not flush buffer.* - See https'
@@ -389,7 +389,7 @@ def test_auto_flush_on_closed_socket(self):
     def test_dont_auto_flush(self):
         msg_counter = 0
         with Server() as server:
-            with qi.Sender('localhost', server.port, auto_flush=0) as sender:
+            with qi.Sender('localhost', server.port, auto_flush=qi.AutoFlush.Disabled) as sender:
                 server.accept()
                 while len(sender) < 32768:  # 32KiB
                     sender.row('tbl1', symbols={'sym1': 'val1'})
@@ -433,7 +433,7 @@ def test_dataframe_auto_flush(self):
         with Server() as server:
             # An auto-flush size of 20 bytes is enough to auto-flush the first
             # row, but not the second.
-            with qi.Sender('localhost', server.port, auto_flush=20) as sender:
+            with qi.Sender('localhost', server.port, auto_flush=qi.AutoFlush.ByteCount(20)) as sender:
                 server.accept()
                 df = pd.DataFrame({'a': [100000, 2], 'b': [3.0, 4.0]})
                 sender.dataframe(df, table_name='tbl1')