feat: (Python) Add async context manager

bindings/python/example/example.py

@@ -667,6 +667,26 @@ async def main():
     except Exception as e:
         print(f"Error during projection: {e}")
 
+
+    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,
+                    "age": 25,
+                    "birth_date": date(2000, 1, 1),
+                    "check_in_time": dt_time(12, 0, 0),
+                    "created_at": datetime(2024, 1, 1, 12, 0, 0),
+                    "updated_at": datetime(2024, 1, 1, 12, 0, 0),
+                    "salary": Decimal("100.00"),
+                }
+            )
+            # auto-flushes on exit
+
     # Demo: Drop tables
     print("\n--- Testing drop_table() ---")
     try:
@@ -933,8 +953,10 @@ async def main():
         print(f"Error with partitioned KV table: {e}")
         traceback.print_exc()
 
+
+
     # Close connection
-    conn.close()
+    await conn.close()
     print("\nConnection closed")
 
 

bindings/python/fluss/__init__.pyi

@@ -245,14 +245,21 @@ class FlussConnection:
     async def create(config: Config) -> FlussConnection: ...
     def get_admin(self) -> FlussAdmin: ...
     async def get_table(self, table_path: TablePath) -> FlussTable: ...
-    def close(self) -> None: ...
+    async def close(self) -> None: ...
     def __enter__(self) -> FlussConnection: ...
     def __exit__(
         self,
         exc_type: Optional[type],
         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,27 @@ 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 exit, the writer is automatically flushed to ensure
+        all pending records are sent and acknowledged.
+        """
+        ...
     def __repr__(self) -> str: ...
 
 class UpsertWriter:
@@ -644,6 +672,27 @@ 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 exit, the writer is automatically flushed to ensure
+        all pending records are sent and acknowledged.
+        """
+        ...
     def __repr__(self) -> str: ...
 
 
@@ -807,6 +856,8 @@ class LogScanner:
 
         You must call subscribe(), subscribe_buckets(), or subscribe_partition() first.
         """
+        ...
+
     def __repr__(self) -> str: ...
     def __aiter__(self) -> AsyncIterator[Union[ScanRecord, RecordBatch]]: ...
 

bindings/python/src/connection.rs

@@ -18,6 +18,7 @@
 use crate::*;
 use pyo3_async_runtimes::tokio::future_into_py;
 use std::sync::Arc;
+use std::time::Duration;
 
 /// Connection to a Fluss cluster
 #[pyclass]
@@ -82,9 +83,19 @@ 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.
+    fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let inner = self.inner.clone();
+
+        future_into_py(py, async move {
+            inner
+                .close(Duration::MAX)
+                .await
+                .map_err(|e| FlussError::from_core_error(&e))
+        })
     }
 
     // Enter the runtime context (for 'with' statement)
@@ -100,10 +111,39 @@ 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();
+        let is_exc_none = exc_type.as_ref().is_none_or(|e| e.is_none());
+        future_into_py(py, async move {
+            let res = inner.close(Duration::MAX).await;
+            if let Err(e) = res {
+                if is_exc_none {
+                    return Err(FlussError::from_core_error(&e));
+                }
+            }
+            Ok(false)
+        })
+    }
+
     fn __repr__(&self) -> String {
         "FlussConnection()".to_string()
     }

bindings/python/src/table.rs

@@ -989,6 +989,35 @@ 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 exit, the writer is automatically flushed.
+    #[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();
+        let is_exc_none = exc_type.as_ref().is_none_or(|e| e.is_none());
+        future_into_py(py, async move {
+            let res = inner.flush().await;
+            if let Err(e) = res {
+                if is_exc_none {
+                    return Err(FlussError::from_core_error(&e));
+                }
+            }
+            Ok(false)
+        })
+    }
+
     fn __repr__(&self) -> String {
         "AppendWriter()".to_string()
     }

bindings/python/src/upsert.rs

@@ -108,6 +108,35 @@ 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 exit, the writer is automatically flushed.
+    #[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 writer = self.writer.clone();
+        let is_exc_none = exc_type.as_ref().is_none_or(|e| e.is_none());
+        future_into_py(py, async move {
+            let res = writer.flush().await;
+            if let Err(e) = res {
+                if is_exc_none {
+                    return Err(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
@@ -151,3 +151,34 @@ def plaintext_bootstrap_servers(fluss_cluster):
 @pytest_asyncio.fixture
 async def admin(connection):
     return connection.get_admin()
+
+
+@pytest_asyncio.fixture
+async def wait_for_table_ready(admin):
+    """
+    Fixture that returns a helper function to wait for a table or partition to be ready.
+    """
+    async def _wait(table_path, timeout=15, partition_name=None):
+        start_time = time.monotonic()
+        while time.monotonic() - start_time < timeout:
+            try:
+                if partition_name:
+                    await admin.list_partition_offsets(
+                        table_path, partition_name, [0], fluss.OffsetSpec.earliest()
+                    )
+                else:
+                    await admin.list_offsets(table_path, [0], fluss.OffsetSpec.earliest())
+                return
+            except (fluss.FlussError, Exception) as e:
+                # Catch "No leader found" or other errors that indicate the table/partition is still initializing
+                err_msg = str(e)
+                if any(msg in err_msg for msg in ["No leader found", "Table not ready", "Metadata not ready", "not leader or follower"]):
+                    await asyncio.sleep(1)
+                    continue
+                raise
+        raise TimeoutError(
+            f"Table/Partition {table_path} ({partition_name or 'standard'}) "
+            f"did not become ready within {timeout}s"
+        )
+
+    return _wait