Updated to support the latest schema