fix: Apply SQLite optimizations to the custom connection_string in SqlStorageClient

src/crawlee/storage_clients/_sql/_storage_client.py

@@ -72,8 +72,10 @@ def __init__(
         self._initialized = False
         self.session_maker: None | async_sessionmaker[AsyncSession] = None
 
-        # Flag needed to apply optimizations only for default database
-        self._default_flag = self._engine is None and self._connection_string is None
+        # Flag needed to apply optimizations only for SQLite database
+        self._optimization_flag = self._engine is None and (
+            self._connection_string is None or self._connection_string.startswith('sqlite')
+        )
         self._dialect_name: str | None = None
 
         # Call the notification only once
@@ -129,7 +131,7 @@ async def initialize(self, configuration: Configuration) -> None:
                 # This is likely an attempt to create a database from several parallel processes.
                 try:
                     # Set SQLite pragmas for performance and consistency
-                    if self._default_flag:
+                    if self._optimization_flag:
                         await conn.execute(text('PRAGMA journal_mode=WAL'))  # Better concurrency
                         await conn.execute(text('PRAGMA synchronous=NORMAL'))  # Balanced safety/speed
                         await conn.execute(text('PRAGMA cache_size=100000'))  # 100MB cache

tests/unit/storage_clients/_sql/test_sql_storage_client.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from crawlee.configuration import Configuration
+from crawlee.storage_clients import SqlStorageClient
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+async def test_sqlite_wal_mode_with_default_connection(tmp_path: Path) -> None:
+    """Test that WAL mode is applied for the default SQLite connection."""
+    configuration = Configuration(storage_dir=str(tmp_path))
+
+    async with SqlStorageClient() as storage_client:
+        await storage_client.initialize(configuration)
+
+        async with storage_client.engine.begin() as conn:
+            result = await conn.execute(text('PRAGMA journal_mode'))
+            assert result.scalar() == 'wal'
+
+
+async def test_sqlite_wal_mode_with_connection_string(tmp_path: Path) -> None:
+    """Test that WAL mode is applied when using a custom SQLite connection string."""
+    db_path = tmp_path / 'test.db'
+    configuration = Configuration(storage_dir=str(tmp_path))
+
+    async with SqlStorageClient(connection_string=f'sqlite+aiosqlite:///{db_path}') as storage_client:
+        await storage_client.initialize(configuration)
+
+        async with storage_client.engine.begin() as conn:
+            result = await conn.execute(text('PRAGMA journal_mode'))
+            assert result.scalar() == 'wal'
+
+
+async def test_sqlite_wal_mode_not_applied_with_custom_engine(tmp_path: Path) -> None:
+    """Test that WAL mode is not applied when using a user-provided engine."""
+    db_path = tmp_path / 'test.db'
+    configuration = Configuration(storage_dir=str(tmp_path))
+    engine = create_async_engine(f'sqlite+aiosqlite:///{db_path}', future=True)
+
+    async with SqlStorageClient(engine=engine) as storage_client:
+        await storage_client.initialize(configuration)
+
+        async with engine.begin() as conn:
+            result = await conn.execute(text('PRAGMA journal_mode'))
+            assert result.scalar() != 'wal'