Extension Attributes – full implementation with batch loading, REST API, unit tests

Extension Attributes are Magento 2’s mechanism for extending existing entities (Product, Order, Customer) without modifying core tables. They participate in the REST API automatically and can carry any data type. Most tutorials show a simple string attribute. I show a complete implementation with batch loading (fixing N+1), REST API exposure with serialisation, and unit tests.

Extension Attributes vs Custom Attributes

Complete example: Order with supplier data

<!-- etc/extension_attributes.xml -->
<?xml version="1.0"?>
<config xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
        xsi:noNamespaceSchemaLocation="urn:magento:framework:Api/etc/extension_attributes.xsd">

    <extension_attributes for="Magento\Sales\Api\Data\OrderInterface">
        <attribute code="supplier_data"
                   type="Vendor\Module\Api\Data\SupplierDataInterface"/>
    </extension_attributes>

</config>

<?php

declare(strict_types=1);

namespace Vendor\Module\Api\Data;

interface SupplierDataInterface
{
    public function getOrderId(): int;
    public function getSupplierId(): ?int;
    public function setOrderId(int $orderId): static;
    public function setSupplierId(?int $supplierId): static;
    public function getSupplierReference(): ?string;
    public function setSupplierReference(?string $ref): static;
    public function getExportedAt(): ?string;
    public function setExportedAt(?string $date): static;
}

Plugin with batch loading – fixing N+1

<?php

declare(strict_types=1);

namespace Vendor\Module\Plugin;

use Magento\Sales\Api\Data\OrderInterface;
use Magento\Sales\Api\Data\OrderSearchResultInterface;

// Plugin on OrderRepository - loads supplier data for all orders at once
class OrderExtensionAttributesPlugin
{
    public function __construct(
        private \Vendor\Module\Model\ResourceModel\SupplierData $supplierDataResource,
        private \Vendor\Module\Api\Data\SupplierDataInterfaceFactory $supplierDataFactory,
        private \Magento\Sales\Api\Data\OrderExtensionFactory $orderExtensionFactory
    ) {}

    // Single order - afterGet
    public function afterGet(
        \Magento\Sales\Api\OrderRepositoryInterface $subject,
        OrderInterface $order
    ): OrderInterface {
        $this->loadForOrders([$order]);
        return $order;
    }

    // Collection - afterGetList (batch load avoids N+1)
    public function afterGetList(
        \Magento\Sales\Api\OrderRepositoryInterface $subject,
        OrderSearchResultInterface $searchResult
    ): OrderSearchResultInterface {
        $orders = $searchResult->getItems();
        if (!empty($orders)) {
            $this->loadForOrders($orders);
        }
        return $searchResult;
    }

    /** @param OrderInterface[] $orders */
    private function loadForOrders(array $orders): void
    {
        $orderIds = array_map(fn($o) => (int)$o->getEntityId(), $orders);

        // ONE query for all orders - no N+1
        $supplierDataByOrderId = $this->supplierDataResource->loadBatch($orderIds);

        foreach ($orders as $order) {
            $orderId    = (int)$order->getEntityId();
            $rawData    = $supplierDataByOrderId[$orderId] ?? null;

            $supplierData = $this->supplierDataFactory->create();
            $supplierData->setOrderId($orderId);

            if ($rawData !== null) {
                $supplierData->setSupplierId($rawData['supplier_id'] ?? null);
                $supplierData->setSupplierReference($rawData['supplier_reference'] ?? null);
                $supplierData->setExportedAt($rawData['exported_at'] ?? null);
            }

            // Attach to order extension attributes
            $extension = $order->getExtensionAttributes() ?? $this->orderExtensionFactory->create();
            $extension->setSupplierData($supplierData);
            $order->setExtensionAttributes($extension);
        }
    }
}

// The ResourceModel with batch loading
class SupplierDataResourceModel extends \Magento\Framework\Model\ResourceModel\Db\AbstractDb
{
    protected function _construct(): void
    {
        $this->_init('vendor_order_supplier_data', 'order_id');
    }

    public function loadBatch(array $orderIds): array
    {
        if (empty($orderIds)) return [];

        $connection = $this->getConnection();
        $rows = $connection->fetchAll(
            $connection->select()
                ->from($this->getMainTable())
                ->where('order_id IN (?)', $orderIds)
        );

        $indexed = [];
        foreach ($rows as $row) {
            $indexed[$row['order_id']] = $row;
        }
        return $indexed;
    }
}

Saving extension attributes

<?php

// Plugin on afterSave to persist extension attribute
public function afterSave(
    \Magento\Sales\Api\OrderRepositoryInterface $subject,
    OrderInterface $order
): OrderInterface {
    $supplierData = $order->getExtensionAttributes()?->getSupplierData();

    if ($supplierData !== null) {
        $connection = $this->resourceConnection->getConnection();
        $table = $this->resourceConnection->getTableName('vendor_order_supplier_data');

        $data = [
            'order_id'           => $order->getEntityId(),
            'supplier_id'        => $supplierData->getSupplierId(),
            'supplier_reference' => $supplierData->getSupplierReference(),
            'exported_at'        => $supplierData->getExportedAt(),
        ];

        $connection->insertOnDuplicate($table, $data);
    }

    return $order;
}

Unit tests

<?php

class OrderExtensionAttributesPluginTest extends TestCase
{
    public function testAfterGetLoadsSupplierData(): void
    {
        $orderId = 42;
        $order   = $this->createMock(OrderInterface::class);
        $order->method('getEntityId')->willReturn($orderId);

        $supplierData = $this->createMock(SupplierDataInterface::class);
        $supplierData->method('getSupplierId')->willReturn(7);

        $resource = $this->createMock(SupplierDataResourceModel::class);
        $resource->method('loadBatch')
            ->with([$orderId])
            ->willReturn([$orderId => ['supplier_id' => 7, 'supplier_reference' => 'REF-001']]);

        $factory = $this->createMock(SupplierDataInterfaceFactory::class);
        $factory->method('create')->willReturn($supplierData);

        $extensionFactory = $this->createMock(OrderExtensionFactory::class);
        $extension = $this->createMock(OrderExtensionInterface::class);
        $extensionFactory->method('create')->willReturn($extension);
        $order->method('getExtensionAttributes')->willReturn(null);

        $extension->expects($this->once())->method('setSupplierData')->with($supplierData);
        $order->expects($this->once())->method('setExtensionAttributes')->with($extension);

        $plugin = new OrderExtensionAttributesPlugin($resource, $factory, $extensionFactory);
        $result = $plugin->afterGet($this->createMock(OrderRepositoryInterface::class), $order);

        $this->assertSame($order, $result);
    }
}

Summary

Extension Attributes are the proper Magento 2 way to extend entities. The critical detail is batch loading – always load extension attributes for all items in a collection with one query, never in a loop. The plugin-on-getList + loadBatch pattern is the established standard. REST API serialisation is automatic once the factory and interface are registered. Unit testing the plugin is straightforward because all dependencies are mockable interfaces.