Some fixes

Author: shefeek-jinnah

src/metadata_writer_sqlite.rs

@@ -1,11 +1,33 @@
 //! SQLite metadata writer for DuckLake catalogs.
+//!
+//! This module provides [`SqliteMetadataWriter`], an implementation of the
+//! [`MetadataWriter`] trait using SQLite as the catalog backend.
+//!
+//! # Runtime Requirements
+//!
+//! This implementation uses `tokio::task::block_in_place` internally, which requires
+//! a **multi-threaded Tokio runtime**. Using `#[tokio::test]` without
+//! `flavor = "multi_thread"` will panic.
+//!
+//! # Example
+//!
+//! ```ignore
+//! use datafusion_ducklake::SqliteMetadataWriter;
+//!
+//! // Create and initialize a new catalog
+//! let writer = SqliteMetadataWriter::new_with_init("sqlite:catalog.db?mode=rwc").await?;
+//! writer.set_data_path("/path/to/data")?;
+//! ```
 
 use crate::Result;
 use crate::metadata_provider::block_on;
 use crate::metadata_writer::{ColumnDef, DataFileInfo, MetadataWriter};
 use sqlx::Row;
 use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
 
+/// Default maximum number of connections in the pool.
+const DEFAULT_MAX_CONNECTIONS: u32 = 5;
+
 /// SQL to create DuckLake schema tables
 const SQL_CREATE_SCHEMA: &str = r#"
 CREATE TABLE IF NOT EXISTS ducklake_metadata (
@@ -78,6 +100,19 @@ CREATE TABLE IF NOT EXISTS ducklake_delete_file (
 "#;
 
 /// SQLite-based metadata writer for DuckLake catalogs.
+///
+/// This writer manages DuckLake catalog metadata stored in a SQLite database,
+/// including snapshots, schemas, tables, columns, and data file registrations.
+///
+/// # Thread Safety
+///
+/// This type is `Clone`, `Send`, and `Sync`. The underlying connection pool
+/// handles concurrent access safely.
+///
+/// # Connection Pool
+///
+/// Uses a connection pool with configurable size (default: 5 connections).
+/// For embedded single-threaded use cases, consider using `with_max_connections(1)`.
 #[derive(Debug, Clone)]
 pub struct SqliteMetadataWriter {
     pool: SqlitePool,
@@ -86,11 +121,43 @@ pub struct SqliteMetadataWriter {
 impl SqliteMetadataWriter {
     /// Creates a new writer for a DuckLake catalog.
     ///
-    /// Connection string format: `sqlite:///path/to/catalog.db?mode=rwc`
-    /// Use `?mode=rwc` to create the file if it doesn't exist.
+    /// # Connection String Format
+    ///
+    /// - `sqlite:path/to/catalog.db` - Open existing database
+    /// - `sqlite:path/to/catalog.db?mode=rwc` - Create if not exists
+    /// - `sqlite::memory:` - In-memory database (for testing)
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// let writer = SqliteMetadataWriter::new("sqlite:catalog.db?mode=rwc").await?;
+    /// ```
     pub async fn new(connection_string: &str) -> Result<Self> {
+        Self::with_max_connections(connection_string, DEFAULT_MAX_CONNECTIONS).await
+    }
+
+    /// Creates a new writer with a custom connection pool size.
+    ///
+    /// # Arguments
+    ///
+    /// * `connection_string` - SQLite connection string
+    /// * `max_connections` - Maximum number of connections in the pool
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// // Single connection for embedded use
+    /// let writer = SqliteMetadataWriter::with_max_connections(
+    ///     "sqlite:catalog.db?mode=rwc",
+    ///     1
+    /// ).await?;
+    /// ```
+    pub async fn with_max_connections(
+        connection_string: &str,
+        max_connections: u32,
+    ) -> Result<Self> {
         let pool = SqlitePoolOptions::new()
-            .max_connections(5)
+            .max_connections(max_connections)
             .connect(connection_string)
             .await?;
 
@@ -99,7 +166,17 @@ impl SqliteMetadataWriter {
         })
     }
 
-    /// Creates a new writer and initializes the schema if needed.
+    /// Creates a new writer and initializes the DuckLake schema tables.
+    ///
+    /// This is a convenience method that combines `new()` and `initialize_schema()`.
+    /// Use this when creating a new catalog from scratch.
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// let writer = SqliteMetadataWriter::new_with_init("sqlite:catalog.db?mode=rwc").await?;
+    /// writer.set_data_path("/data/warehouse")?;
+    /// ```
     pub async fn new_with_init(connection_string: &str) -> Result<Self> {
         let writer = Self::new(connection_string).await?;
         writer.initialize_schema()?;

src/table_writer.rs

@@ -317,7 +317,7 @@ impl DuckLakeTableWriter {
 
 /// A streaming write session for writing batches incrementally to a DuckLake table.
 ///
-/// Created by `DuckLakeTableWriter::begin_write()` or `begin_write_to_path()`.
+/// Created by [`DuckLakeTableWriter::begin_write()`] or [`DuckLakeTableWriter::begin_write_to_path()`].
 ///
 /// The session holds the Parquet writer and catalog metadata. Batches are written
 /// incrementally via `write_batch()`, and the file is registered in the catalog
@@ -329,8 +329,29 @@ impl DuckLakeTableWriter {
 /// begin_write() → write_batch()* → finish()
 /// ```
 ///
-/// If the session is dropped without calling `finish()`, the file is written but
-/// NOT registered in the catalog (effectively discarded).
+/// # Error Handling and Cleanup
+///
+/// - **Dropped without `finish()`**: The Parquet file may be partially written but is
+///   NOT registered in the catalog. The orphaned file remains on disk and should be
+///   cleaned up by the caller or a separate garbage collection process.
+///
+/// - **Error during `write_batch()`**: The session remains valid and you can continue
+///   writing. The Parquet file may contain partial data.
+///
+/// - **Error during `finish()`**: The Parquet file is closed but may not be registered
+///   in the catalog. The file remains on disk.
+///
+/// For robust error handling in production, consider:
+/// 1. Wrapping writes in a try block
+/// 2. Implementing cleanup logic for orphaned files
+/// 3. Using the file path from `file_path()` to clean up on error
+///
+/// # Atomicity
+///
+/// The catalog registration in `finish()` is the commit point. Until `finish()`
+/// completes successfully, the data is not visible to readers. However, the
+/// snapshot, schema, table, and columns are created during `begin_write()`,
+/// so a failed write may leave empty catalog entries.
 #[derive(Debug)]
 pub struct TableWriteSession {
     metadata: Arc<dyn MetadataWriter>,