feat: (Python) Add async context manager

bindings/python/example/example.py

@@ -933,8 +933,15 @@ async def main():
         print(f"Error with partitioned KV table: {e}")
         traceback.print_exc()
 
+
+    print("\n--- New: async context manager demo ---")
+    async with await fluss.FlussConnection.create(config) as demo_conn:
+        demo_table = await demo_conn.get_table(table_path)
+        async with demo_table.new_append().create_writer() as writer:
+            writer.append({"id": 1, "name": "demo", "score": 1.0})
+            # auto-flushes on exit
     # Close connection
-    conn.close()
+    await conn.close()
     print("\nConnection closed")
 
 

bindings/python/fluss/__init__.pyi

@@ -253,6 +253,13 @@ class FlussConnection:
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> bool: ...
+    async def __aenter__(self) -> FlussConnection: ...
+    async def __aexit__(
+        self,
+        exc_type: Optional[type],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool: ...
     def __repr__(self) -> str: ...
 
 class ServerNode:
@@ -611,6 +618,37 @@ class AppendWriter:
     def write_arrow_batch(self, batch: pa.RecordBatch) -> WriteResultHandle: ...
     def write_pandas(self, df: pd.DataFrame) -> None: ...
     async def flush(self) -> None: ...
+    async def __aenter__(self) -> AppendWriter:
+        """
+        Enter the async context manager.
+
+        Returns:
+            The AppendWriter instance.
+        """
+        ...
+    async def __aexit__(
+        self,
+        exc_type: Optional[type],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool:
+        """
+        Exit the async context manager.
+
+        On successful exit, the writer is automatically flushed to ensure
+        all pending records are sent and acknowledged.
+
+        Note on Exceptions:
+            If an exception occurs inside the `async with` block, `flush()` is
+            bypassed to return control to the event loop immediately. However,
+            any records already passed to `append()` prior to the exception
+            reside in a shared background buffer and will still be transmitted
+            to the server.
+
+            To achieve true atomicity, buffer your records in a Python list and
+            write them in a single batch at the end of your logic.
+        """
+        ...
     def __repr__(self) -> str: ...
 
 class UpsertWriter:
@@ -644,6 +682,37 @@ class UpsertWriter:
     async def flush(self) -> None:
         """Flush all pending upsert/delete operations to the server."""
         ...
+    async def __aenter__(self) -> UpsertWriter:
+        """
+        Enter the async context manager.
+
+        Returns:
+            The UpsertWriter instance.
+        """
+        ...
+    async def __aexit__(
+        self,
+        exc_type: Optional[type],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool:
+        """
+        Exit the async context manager.
+
+        On successful exit, the writer is automatically flushed to ensure
+        all pending records are sent and acknowledged.
+
+        Note on Exceptions:
+            If an exception occurs inside the `async with` block, `flush()` is
+            bypassed to return control to the event loop immediately. However,
+            any records already passed to `upsert()` or `delete()` prior to the
+            exception reside in a shared background buffer and will still be
+            transmitted to the server.
+
+            To achieve true atomicity, buffer your records in a Python list and
+            write them in a single batch at the end of your logic.
+        """
+        ...
     def __repr__(self) -> str: ...
 
 
@@ -807,6 +876,15 @@ class LogScanner:
 
         You must call subscribe(), subscribe_buckets(), or subscribe_partition() first.
         """
+        ...
+    def close(self) -> None: ...
+    async def __aenter__(self) -> LogScanner: ...
+    async def __aexit__(
+        self,
+        exc_type: Optional[type],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool: ...
     def __repr__(self) -> str: ...
     def __aiter__(self) -> AsyncIterator[Union[ScanRecord, RecordBatch]]: ...
 

bindings/python/src/connection.rs

@@ -18,6 +18,9 @@
 use crate::*;
 use pyo3_async_runtimes::tokio::future_into_py;
 use std::sync::Arc;
+use std::time::Duration;
+
+const DEFAULT_CLOSE_TIMEOUT_SECS: u64 = 30;
 
 /// Connection to a Fluss cluster
 #[pyclass]
@@ -82,9 +85,27 @@ impl FlussConnection {
         })
     }
 
-    // Close the connection
-    fn close(&mut self) -> PyResult<()> {
-        Ok(())
+    /// Close the connection (async).
+    ///
+    /// Gracefully shuts down the connection by draining any pending write batches.
+    /// This method is awaitable.
+    ///
+    /// Args:
+    ///     timeout_ms: The timeout in milliseconds to wait for the graceful drain.
+    ///                 If not provided, defaults to 30 seconds.
+    #[pyo3(signature = (timeout_ms=None))]
+    fn close<'py>(&self, py: Python<'py>, timeout_ms: Option<u64>) -> PyResult<Bound<'py, PyAny>> {
+        let inner = self.inner.clone();
+        let timeout = timeout_ms
+            .map(Duration::from_millis)
+            .unwrap_or_else(|| Duration::from_secs(DEFAULT_CLOSE_TIMEOUT_SECS));
+
+        future_into_py(py, async move {
+            inner
+                .close(timeout)
+                .await
+                .map_err(|e| FlussError::from_core_error(&e))
+        })
     }
 
     // Enter the runtime context (for 'with' statement)
@@ -100,10 +121,36 @@ impl FlussConnection {
         _exc_value: Option<Bound<'_, PyAny>>,
         _traceback: Option<Bound<'_, PyAny>>,
     ) -> PyResult<bool> {
-        self.close()?;
+        // Sync exit cannot await the graceful drain, so it's a no-op here.
+        // Users should use 'async with' for graceful shutdown.
         Ok(false)
     }
 
+    // Enter the async runtime context (for 'async with' statement)
+    fn __aenter__<'py>(slf: PyRef<'py, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let py_slf = slf.into_pyobject(py)?.unbind();
+        future_into_py(py, async move { Ok(py_slf) })
+    }
+
+    // Exit the async runtime context (for 'async with' statement)
+    #[pyo3(signature = (_exc_type=None, _exc_value=None, _traceback=None))]
+    fn __aexit__<'py>(
+        &self,
+        py: Python<'py>,
+        _exc_type: Option<Bound<'py, PyAny>>,
+        _exc_value: Option<Bound<'py, PyAny>>,
+        _traceback: Option<Bound<'py, PyAny>>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let inner = self.inner.clone();
+        future_into_py(py, async move {
+            inner
+                .close(Duration::from_secs(DEFAULT_CLOSE_TIMEOUT_SECS))
+                .await
+                .map_err(|e| FlussError::from_core_error(&e))?;
+            Ok(false)
+        })
+    }
+
     fn __repr__(&self) -> String {
         "FlussConnection()".to_string()
     }

bindings/python/src/table.rs

@@ -989,6 +989,37 @@ impl AppendWriter {
         })
     }
 
+    // Enter the async runtime context (for 'async with' statement)
+    fn __aenter__<'py>(slf: PyRef<'py, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let py_slf = slf.into_pyobject(py)?.unbind();
+        future_into_py(py, async move { Ok(py_slf) })
+    }
+
+    // Exit the async runtime context (for 'async with' statement)
+    /// On successful exit, the writer is automatically flushed.
+    /// If an exception occurs, the flush is skipped to allow immediate error
+    /// propagation, though pending records may still be sent in the background.
+    #[pyo3(signature = (exc_type=None, _exc_value=None, _traceback=None))]
+    fn __aexit__<'py>(
+        &self,
+        py: Python<'py>,
+        exc_type: Option<Bound<'py, PyAny>>,
+        _exc_value: Option<Bound<'py, PyAny>>,
+        _traceback: Option<Bound<'py, PyAny>>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let has_error = exc_type.is_some();
+        let inner = self.inner.clone();
+        future_into_py(py, async move {
+            if !has_error {
+                inner
+                    .flush()
+                    .await
+                    .map_err(|e| FlussError::from_core_error(&e))?;
+            }
+            Ok(false)
+        })
+    }
+
     fn __repr__(&self) -> String {
         "AppendWriter()".to_string()
     }

bindings/python/src/upsert.rs

@@ -108,6 +108,37 @@ impl UpsertWriter {
         })
     }
 
+    // Enter the async runtime context (for 'async with' statement)
+    fn __aenter__<'py>(slf: PyRef<'py, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let py_slf = slf.into_pyobject(py)?.unbind();
+        future_into_py(py, async move { Ok(py_slf) })
+    }
+
+    // Exit the async runtime context (for 'async with' statement)
+    /// On successful exit, the writer is automatically flushed.
+    /// If an exception occurs, the flush is skipped to allow immediate error
+    /// propagation, though pending records may still be sent in the background.
+    #[pyo3(signature = (exc_type=None, _exc_value=None, _traceback=None))]
+    fn __aexit__<'py>(
+        &self,
+        py: Python<'py>,
+        exc_type: Option<Bound<'py, PyAny>>,
+        _exc_value: Option<Bound<'py, PyAny>>,
+        _traceback: Option<Bound<'py, PyAny>>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let has_error = exc_type.is_some();
+        let writer = self.writer.clone();
+        future_into_py(py, async move {
+            if !has_error {
+                writer
+                    .flush()
+                    .await
+                    .map_err(|e| FlussError::from_core_error(&e))?;
+            }
+            Ok(false)
+        })
+    }
+
     fn __repr__(&self) -> String {
         "UpsertWriter()".to_string()
     }

bindings/python/test/conftest.py

@@ -96,7 +96,7 @@ async def _connect(bootstrap_servers):
             nodes = await admin.get_server_nodes()
             if any(n.server_type == "TabletServer" for n in nodes):
                 return conn
-            conn.close()
+            await conn.close()
             last_err = RuntimeError("No TabletServer available yet")
         except Exception as e:
             last_err = e