chore: import upstream snapshot with attribution

backend/tests/data/conftest.py

@@ -0,0 +1,21 @@
+from unittest.mock import AsyncMock
+
+import pytest
+
+from chainlit.data.storage_clients.base import BaseStorageClient
+from chainlit.user import User
+
+
+@pytest.fixture
+def mock_storage_client():
+    mock_client = AsyncMock(spec=BaseStorageClient)
+    mock_client.upload_file.return_value = {
+        "url": "https://example.com/test.txt",
+        "object_key": "test_user/test_element/test.txt",
+    }
+    return mock_client
+
+
+@pytest.fixture
+def test_user() -> User:
+    return User(identifier="test_user_identifier", metadata={})

backend/tests/data/storage_clients/test_gcs.py

@@ -0,0 +1,337 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from chainlit.data.storage_clients.base import storage_expiry_time
+from chainlit.data.storage_clients.gcs import GCSStorageClient
+
+
+@pytest.fixture
+def mock_gcs_client():
+    """Create a mock Google Cloud Storage client."""
+    # First mock the service_account
+    with patch(
+        "chainlit.data.storage_clients.gcs.service_account"
+    ) as mock_service_account:
+        mock_credentials = MagicMock()
+        mock_service_account.Credentials.from_service_account_info.return_value = (
+            mock_credentials
+        )
+
+        # Then mock the storage client
+        with patch("chainlit.data.storage_clients.gcs.storage") as mock_storage:
+            mock_client = MagicMock()
+            mock_storage.Client.return_value = mock_client
+            mock_bucket = MagicMock()
+            mock_client.bucket.return_value = mock_bucket
+            mock_blob = MagicMock()
+            mock_bucket.blob.return_value = mock_blob
+
+            yield {
+                "storage": mock_storage,
+                "client": mock_client,
+                "bucket": mock_bucket,
+                "blob": mock_blob,
+                "service_account": mock_service_account,
+                "credentials": mock_credentials,
+            }
+
+
+class TestGCSStorageClient:
+    def test_init(self, mock_gcs_client):
+        """Test initialization of GCS client."""
+        # Remove client assignment, directly call the constructor
+        GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Verify service account credentials were created correctly
+        mock_gcs_client[
+            "service_account"
+        ].Credentials.from_service_account_info.assert_called_once_with(
+            {
+                "type": "service_account",
+                "project_id": "test-project",
+                "private_key": "test-key",
+                "client_email": "test@example.com",
+                "token_uri": "https://oauth2.googleapis.com/token",
+            }
+        )
+
+        # Verify storage client was initialized with credentials
+        mock_gcs_client["storage"].Client.assert_called_once_with(
+            project="test-project", credentials=mock_gcs_client["credentials"]
+        )
+
+        # Verify bucket was retrieved
+        mock_gcs_client["client"].bucket.assert_called_once_with("test-bucket")
+
+    def test_sync_get_read_url(self, mock_gcs_client):
+        """Test generating a signed URL."""
+        # Set up the mock to return a URL
+        mock_gcs_client[
+            "blob"
+        ].generate_signed_url.return_value = "https://signed-url.example.com"
+
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset mocks to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Test the method
+        url = client.sync_get_read_url("test/path/file.txt")
+
+        # Verify correct methods were called
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].generate_signed_url.assert_called_once_with(
+            version="v4", expiration=storage_expiry_time, method="GET"
+        )
+
+        assert url == "https://signed-url.example.com"
+
+    @pytest.mark.asyncio
+    async def test_get_read_url(self, mock_gcs_client):
+        """Test the async wrapper for getting a read URL."""
+        # Set up the mock to return a URL
+        mock_gcs_client[
+            "blob"
+        ].generate_signed_url.return_value = "https://signed-url.example.com"
+
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset mocks to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Test the async method
+        url = await client.get_read_url("test/path/file.txt")
+
+        # Verify correct methods were called
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].generate_signed_url.assert_called_once_with(
+            version="v4", expiration=storage_expiry_time, method="GET"
+        )
+
+        assert url == "https://signed-url.example.com"
+
+    def test_sync_upload_file(self, mock_gcs_client):
+        """Test uploading a file to GCS."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Mock the bucket name property
+        mock_gcs_client["bucket"].name = "test-bucket"
+
+        # Test with binary data
+        binary_data = b"test content"
+        object_key = "test/path/file.txt"
+
+        # Remove the unused result assignment
+        client.sync_upload_file(
+            object_key=object_key, data=binary_data, mime="text/plain", overwrite=True
+        )
+
+        # Check that the blob was called with the correct object key (using assert_any_call instead of assert_called_once_with)
+        mock_gcs_client["bucket"].blob.assert_any_call(object_key)
+        mock_gcs_client["blob"].upload_from_string.assert_called_once_with(
+            binary_data, content_type="text/plain"
+        )
+
+    def test_sync_upload_file_string_data(self, mock_gcs_client):
+        """Test uploading string data to GCS."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+        mock_gcs_client["bucket"].name = "test-bucket"
+
+        # Test with string data that should be encoded
+        string_data = "test content"
+        object_key = "test/path/file.txt"
+
+        # Remove the unused result assignment
+        client.sync_upload_file(
+            object_key=object_key, data=string_data, mime="text/plain", overwrite=True
+        )
+
+        # Check that the correct methods were called
+        mock_gcs_client["bucket"].blob.assert_any_call(object_key)
+        mock_gcs_client["blob"].upload_from_string.assert_called_once_with(
+            b"test content", content_type="text/plain"
+        )
+
+    def test_sync_upload_file_no_overwrite(self, mock_gcs_client):
+        """Test upload with overwrite=False and existing file."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Configure blob to indicate file exists
+        mock_gcs_client["blob"].exists.return_value = True
+
+        with pytest.raises(
+            Exception,
+            match=r"Failed to upload file to GCS: File test/path/existing\.txt already exists and overwrite is False",
+        ):
+            client.sync_upload_file(
+                object_key="test/path/existing.txt",
+                data=b"test content",
+                overwrite=False,
+            )
+
+        mock_gcs_client["blob"].exists.assert_called_once()
+        mock_gcs_client["blob"].upload_from_string.assert_not_called()
+
+    def test_sync_upload_file_error(self, mock_gcs_client):
+        """Test error handling during upload."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Configure upload to throw an exception
+        mock_gcs_client["blob"].upload_from_string.side_effect = ValueError(
+            "Upload failed"
+        )
+
+        with pytest.raises(
+            Exception, match="Failed to upload file to GCS: Upload failed"
+        ):
+            client.sync_upload_file(
+                object_key="test/path/file.txt", data=b"test content"
+            )
+
+    @pytest.mark.asyncio
+    async def test_upload_file(self, mock_gcs_client):
+        """Test the async wrapper for uploading a file."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+        mock_gcs_client["bucket"].name = "test-bucket"
+
+        # Test with binary data
+        binary_data = b"test content"
+        object_key = "test/path/file.txt"
+
+        # Remove the unused result assignment
+        await client.upload_file(
+            object_key=object_key, data=binary_data, mime="text/plain", overwrite=True
+        )
+
+        # Check that the correct methods were called
+        mock_gcs_client["bucket"].blob.assert_any_call(object_key)
+        mock_gcs_client["blob"].upload_from_string.assert_called_once_with(
+            binary_data, content_type="text/plain"
+        )
+
+    def test_sync_delete_file(self, mock_gcs_client):
+        """Test deleting a file from GCS."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Test successful delete
+        result = client.sync_delete_file("test/path/file.txt")
+
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].delete.assert_called_once()
+        assert result is True
+
+    def test_sync_delete_file_error(self, mock_gcs_client):
+        """Test error handling during file deletion."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Configure delete to throw an exception
+        mock_gcs_client["blob"].delete.side_effect = ValueError("Delete failed")
+
+        # The method should catch the exception and return False
+        result = client.sync_delete_file("test/path/file.txt")
+
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].delete.assert_called_once()
+        assert result is False
+
+    @pytest.mark.asyncio
+    async def test_delete_file(self, mock_gcs_client):
+        """Test the async wrapper for deleting a file."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Test successful delete
+        result = await client.delete_file("test/path/file.txt")
+
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].delete.assert_called_once()
+        assert result is True

backend/tests/data/storage_clients/test_s3.py

@@ -0,0 +1,47 @@
+import os
+
+import boto3  # type: ignore
+import pytest
+from moto import mock_aws
+
+from chainlit.data.storage_clients.s3 import S3StorageClient
+
+
+# Fixtures for setting up the DynamoDB table
+@pytest.fixture
+def aws_credentials():
+    """Mocked AWS Credentials for moto."""
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
+
+
+@pytest.fixture
+def s3_mock(aws_credentials):
+    """Moto mock S3 setup."""
+    with mock_aws():
+        s3 = boto3.client("s3", region_name="us-east-1")
+        # Create a mock bucket
+        s3.create_bucket(Bucket="my-test-bucket")
+        yield s3
+
+
+@pytest.mark.asyncio
+async def test_upload_file(s3_mock):
+    # Initialize the S3StorageClient with the mock bucket
+    client = S3StorageClient(bucket="my-test-bucket")
+
+    # Call the upload_file method and await the result
+    result = await client.upload_file(
+        object_key="test.txt", data="This is a test file", mime="text/plain"
+    )
+
+    # Assert that the file upload returned the correct URL
+    assert result["object_key"] == "test.txt"
+    assert result["url"] == "https://my-test-bucket.s3.amazonaws.com/test.txt"
+
+    # Verify that the file exists in the mock S3
+    response = s3_mock.get_object(Bucket="my-test-bucket", Key="test.txt")
+    assert response["Body"].read().decode() == "This is a test file"

backend/tests/data/test_chainlit_data_layer.py

@@ -0,0 +1,224 @@
+import json
+from unittest.mock import AsyncMock
+
+import pytest
+
+from chainlit.data.chainlit_data_layer import ChainlitDataLayer
+
+
+@pytest.mark.asyncio
+async def test_update_thread_preserves_metadata_when_none_existing_thread():
+    """Test that update_thread preserves existing metadata when metadata=None on an existing thread."""
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    existing_metadata = {"important": "data", "keep": "me"}
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            return [{"metadata": json.dumps(existing_metadata)}]
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    await data_layer.update_thread(thread_id="test-thread-123", name="Updated Name")
+
+    assert data_layer.execute_query.call_count == 2
+
+    update_call = data_layer.execute_query.call_args_list[1]
+    update_params = update_call[0][1]
+
+    metadata_json = None
+    for value in update_params.values():
+        if isinstance(value, str) and value.startswith("{"):
+            try:
+                metadata_json = json.loads(value)
+                break
+            except json.JSONDecodeError:
+                pass
+
+    assert metadata_json == existing_metadata
+
+
+@pytest.mark.asyncio
+async def test_update_thread_noop_skips_upsert_on_existing_thread():
+    """Test that update_thread with no arguments early-returns for an existing thread."""
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            return [{"metadata": json.dumps({"existing": "data"})}]
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    await data_layer.update_thread(thread_id="test-thread-123")
+
+    assert data_layer.execute_query.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_update_thread_new_thread_includes_empty_metadata():
+    """Test that update_thread includes metadata='{}' for a new thread (NOT NULL safe)."""
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            return []
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    await data_layer.update_thread(thread_id="new-thread", name="New Chat")
+
+    assert data_layer.execute_query.call_count == 2
+
+    update_call = data_layer.execute_query.call_args_list[1]
+    update_query = update_call[0][0]
+    update_params = update_call[0][1]
+
+    assert "metadata" in update_query.lower()
+
+    metadata_json = None
+    for value in update_params.values():
+        if isinstance(value, str) and value.startswith("{"):
+            try:
+                metadata_json = json.loads(value)
+                break
+            except json.JSONDecodeError:
+                pass
+
+    assert metadata_json == {}
+
+
+@pytest.mark.asyncio
+async def test_update_thread_merges_metadata_when_provided():
+    """Test that update_thread merges metadata correctly when provided."""
+    # Create a mock data layer
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    # Mock the execute_query method to return existing metadata
+    existing_metadata = {"is_shared": True, "custom_field": "original"}
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            # Return existing thread metadata
+            return [{"metadata": json.dumps(existing_metadata)}]
+        # For the UPDATE/INSERT, return None
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    # Call update_thread with partial metadata update
+    new_metadata = {"custom_field": "updated", "new_field": "added"}
+    await data_layer.update_thread(
+        thread_id="test-thread-123", name="Updated Name", metadata=new_metadata
+    )
+
+    # Verify execute_query was called twice (once for SELECT, once for UPDATE)
+    assert data_layer.execute_query.call_count == 2
+
+    # Get the UPDATE call
+    update_call = data_layer.execute_query.call_args_list[1]
+    update_params = update_call[0][1]
+
+    # The metadata should be merged
+    # Expected: {"is_shared": True, "custom_field": "updated", "new_field": "added"}
+    # Find the JSON metadata in the params
+    metadata_json = None
+    for value in update_params.values():
+        if isinstance(value, str) and value.startswith("{"):
+            try:
+                metadata_json = json.loads(value)
+                break
+            except json.JSONDecodeError:
+                pass
+
+    assert metadata_json is not None
+    assert metadata_json.get("is_shared") is True
+    assert metadata_json.get("custom_field") == "updated"
+    assert metadata_json.get("new_field") == "added"
+
+
+@pytest.mark.asyncio
+async def test_update_thread_deletes_keys_with_none_values():
+    """Test that update_thread deletes keys when value is None."""
+    # Create a mock data layer
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    # Mock the execute_query method to return existing metadata
+    existing_metadata = {
+        "is_shared": True,
+        "to_delete": "will be removed",
+        "keep": "stays",
+    }
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            # Return existing thread metadata
+            return [{"metadata": json.dumps(existing_metadata)}]
+        # For the UPDATE/INSERT, return None
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    # Call update_thread with None value to delete a key
+    new_metadata = {"to_delete": None, "new_field": "added"}
+    await data_layer.update_thread(thread_id="test-thread-123", metadata=new_metadata)
+
+    # Verify execute_query was called twice
+    assert data_layer.execute_query.call_count == 2
+
+    # Get the UPDATE call
+    update_call = data_layer.execute_query.call_args_list[1]
+    update_params = update_call[0][1]
+
+    # The metadata should have deleted "to_delete" key and added "new_field"
+    # Expected: {"is_shared": True, "keep": "stays", "new_field": "added"}
+    metadata_json = None
+    for value in update_params.values():
+        if isinstance(value, str) and value.startswith("{"):
+            try:
+                metadata_json = json.loads(value)
+                break
+            except json.JSONDecodeError:
+                pass
+
+    if metadata_json:
+        # Verify "to_delete" is not in the merged metadata
+        assert "to_delete" not in metadata_json
+        # Verify "new_field" was added
+        assert metadata_json.get("new_field") == "added"
+        # Verify "is_shared" and "keep" are preserved
+        assert metadata_json.get("is_shared") is True
+        assert metadata_json.get("keep") == "stays"
+
+
+@pytest.mark.asyncio
+async def test_create_step_uses_nullif_for_output_and_input():
+    """Empty-string output/input should not overwrite existing content.
+
+    Regression test for https://github.com/Chainlit/chainlit/issues/2789
+    The SQL uses NULLIF(EXCLUDED.output, '') so that an empty string from the
+    initial Step.send() is treated as NULL by COALESCE, preventing it from
+    overwriting non-empty content saved by a subsequent Step.update().
+    """
+    import inspect
+
+    source = inspect.getsource(ChainlitDataLayer.create_step)
+
+    assert "NULLIF(EXCLUDED.output, '')" in source, (
+        "output should use NULLIF to treat empty string as NULL"
+    )
+    assert "NULLIF(EXCLUDED.input, '')" in source, (
+        "input should use NULLIF to treat empty string as NULL"
+    )

backend/tests/data/test_get_data_layer.py

@@ -0,0 +1,18 @@
+from unittest.mock import AsyncMock, Mock
+
+from chainlit.data import get_data_layer
+
+
+async def test_get_data_layer(
+    mock_data_layer: AsyncMock,
+    mock_get_data_layer: Mock,
+):
+    # Check whether the data layer is properly set
+    assert mock_data_layer == get_data_layer()
+
+    mock_get_data_layer.assert_called_once()
+
+    # Getting the data layer again, should not result in additional call
+    assert mock_data_layer == get_data_layer()
+
+    mock_get_data_layer.assert_called_once()

backend/tests/data/test_sql_alchemy.py

@@ -0,0 +1,327 @@
+import json
+import uuid
+from pathlib import Path
+
+import pytest
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from chainlit import User
+from chainlit.data.sql_alchemy import SQLAlchemyDataLayer
+from chainlit.data.storage_clients.base import BaseStorageClient
+from chainlit.element import Text
+
+
+@pytest.fixture
+async def data_layer(mock_storage_client: BaseStorageClient, tmp_path: Path):
+    db_file = tmp_path / "test_db.sqlite"
+    conninfo = f"sqlite+aiosqlite:///{db_file}"
+
+    # Create async engine
+    engine = create_async_engine(conninfo)
+
+    # Execute initialization statements
+    # Ref: https://docs.chainlit.io/data-persistence/custom#sql-alchemy-data-layer
+    async with engine.begin() as conn:
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE users (
+                    "id" UUID PRIMARY KEY,
+                    "identifier" TEXT NOT NULL UNIQUE,
+                    "metadata" JSONB NOT NULL,
+                    "createdAt" TEXT
+                );
+        """
+            )
+        )
+
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS threads (
+                    "id" UUID PRIMARY KEY,
+                    "createdAt" TEXT,
+                    "name" TEXT,
+                    "userId" UUID,
+                    "userIdentifier" TEXT,
+                    "tags" TEXT[],
+                    "metadata" JSONB NOT NULL DEFAULT '{}',
+                    FOREIGN KEY ("userId") REFERENCES users("id") ON DELETE CASCADE
+                );
+        """
+            )
+        )
+
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS steps (
+                    "id" UUID PRIMARY KEY,
+                    "name" TEXT NOT NULL,
+                    "type" TEXT NOT NULL,
+                    "threadId" UUID NOT NULL,
+                    "parentId" UUID,
+                    "disableFeedback" BOOLEAN NOT NULL,
+                    "streaming" BOOLEAN NOT NULL,
+                    "waitForAnswer" BOOLEAN,
+                    "isError" BOOLEAN,
+                    "metadata" JSONB,
+                    "tags" TEXT[],
+                    "input" TEXT,
+                    "output" TEXT,
+                    "createdAt" TEXT,
+                    "start" TEXT,
+                    "end" TEXT,
+                    "generation" JSONB,
+                    "showInput" TEXT,
+                    "language" TEXT,
+                    "indent" INT
+                );
+        """
+            )
+        )
+
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS elements (
+                    "id" UUID PRIMARY KEY,
+                    "threadId" UUID,
+                    "type" TEXT,
+                    "url" TEXT,
+                    "chainlitKey" TEXT,
+                    "name" TEXT NOT NULL,
+                    "display" TEXT,
+                    "objectKey" TEXT,
+                    "size" TEXT,
+                    "page" INT,
+                    "language" TEXT,
+                    "forId" UUID,
+                    "mime" TEXT
+                );
+        """
+            )
+        )
+
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS feedbacks (
+                    "id" UUID PRIMARY KEY,
+                    "forId" UUID NOT NULL,
+                    "threadId" UUID NOT NULL,
+                    "value" INT NOT NULL,
+                    "comment" TEXT
+                );
+        """
+            )
+        )
+
+    # Create SQLAlchemyDataLayer instance
+    data_layer = SQLAlchemyDataLayer(conninfo, storage_provider=mock_storage_client)
+
+    return data_layer
+
+
+async def test_create_and_get_element(
+    mock_chainlit_context, data_layer: SQLAlchemyDataLayer
+):
+    async with mock_chainlit_context:
+        text_element = Text(
+            id=str(uuid.uuid4()),
+            name="test.txt",
+            mime="text/plain",
+            content="test content",
+            for_id="test_step_id",
+        )
+
+        # Needs context because of wrapper in utils.py
+        await data_layer.create_element(text_element)
+
+    retrieved_element = await data_layer.get_element(
+        text_element.thread_id, text_element.id
+    )
+    assert retrieved_element is not None
+    assert retrieved_element["id"] == text_element.id
+    assert retrieved_element["name"] == text_element.name
+    assert retrieved_element["mime"] == text_element.mime
+    # The 'content' field is not part of the ElementDict, so we remove this assertion
+
+
+async def test_get_current_timestamp(data_layer: SQLAlchemyDataLayer):
+    timestamp = await data_layer.get_current_timestamp()
+    assert isinstance(timestamp, str)
+
+
+async def test_get_user(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    fetched_user = await data_layer.get_user(persisted_user.identifier)
+
+    assert fetched_user
+    assert fetched_user.createdAt == persisted_user.createdAt
+    assert fetched_user.id == persisted_user.id
+
+    nonexistent_user = await data_layer.get_user("nonexistent")
+    assert nonexistent_user is None
+
+
+async def test_create_user(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+
+    assert persisted_user
+    assert persisted_user.identifier == test_user.identifier
+    assert persisted_user.createdAt
+    assert persisted_user.id
+
+    # Assert id is valid uuid
+    assert uuid.UUID(persisted_user.id)
+
+
+async def test_update_thread(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    await data_layer.update_thread("test_thread")
+
+
+async def test_get_thread_author(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    await data_layer.update_thread("test_thread", user_id=persisted_user.id)
+    author = await data_layer.get_thread_author("test_thread")
+
+    assert author == persisted_user.identifier
+
+
+async def test_get_thread(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    await data_layer.update_thread("test_thread")
+    result = await data_layer.get_thread("test_thread")
+    assert result is not None
+
+    result = await data_layer.get_thread("nonexisting_thread")
+    assert result is None
+
+
+async def test_delete_thread(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    await data_layer.update_thread("test_thread", "test_user")
+    await data_layer.delete_thread("test_thread")
+    thread = await data_layer.get_thread("test_thread")
+    assert thread is None
+
+
+async def _get_thread_metadata_raw(
+    data_layer: SQLAlchemyDataLayer, thread_id: str
+) -> str | None:
+    result = await data_layer.execute_sql(
+        query='SELECT "metadata" FROM threads WHERE "id" = :id',
+        parameters={"id": thread_id},
+    )
+    if isinstance(result, list) and result:
+        return result[0].get("metadata")
+    return None
+
+
+async def test_update_thread_without_metadata_stores_empty_dict(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread("thread_no_meta")
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_no_meta")
+    assert raw is not None
+    assert json.loads(raw) == {}
+
+
+async def test_update_thread_with_explicit_metadata(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread(
+        "thread_meta", metadata={"key": "value", "count": 42}
+    )
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_meta")
+    assert raw is not None
+    assert json.loads(raw) == {"key": "value", "count": 42}
+
+
+async def test_update_thread_with_empty_metadata(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread("thread_empty_meta", metadata={})
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_empty_meta")
+    assert raw is not None
+    assert json.loads(raw) == {}
+
+
+async def test_update_thread_preserves_metadata_on_noop_call(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread("thread_preserve", metadata={"important": "data"})
+
+    await data_layer.update_thread("thread_preserve")
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_preserve")
+    assert raw is not None
+    assert json.loads(raw) == {"important": "data"}
+
+
+async def test_update_thread_merges_metadata(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread(
+        "thread_merge", metadata={"existing": "old", "keep": "me"}
+    )
+
+    await data_layer.update_thread(
+        "thread_merge", metadata={"existing": "new", "added": "field"}
+    )
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_merge")
+    assert raw is not None
+    merged = json.loads(raw)
+    assert merged == {"existing": "new", "keep": "me", "added": "field"}
+
+
+async def test_update_thread_deletes_metadata_keys_via_none(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread(
+        "thread_delete_key", metadata={"a": 1, "b": 2, "c": 3}
+    )
+
+    await data_layer.update_thread("thread_delete_key", metadata={"b": None})
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_delete_key")
+    assert raw is not None
+    result = json.loads(raw)
+    assert result == {"a": 1, "c": 3}
+
+
+async def test_update_thread_name_update_preserves_metadata(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread("thread_name_meta", metadata={"saved": True})
+
+    await data_layer.update_thread("thread_name_meta", name="New Name")
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_name_meta")
+    assert raw is not None
+    assert json.loads(raw) == {"saved": True}
+
+    result = await data_layer.execute_sql(
+        query='SELECT "name" FROM threads WHERE "id" = :id',
+        parameters={"id": "thread_name_meta"},
+    )
+    assert isinstance(result, list)
+    assert result
+    assert result[0]["name"] == "New Name"