diff --git a/src/Extension/Engine/SiteTreePublishingEngine.php b/src/Extension/Engine/SiteTreePublishingEngine.php
index 1ede852..6f394d6 100644
--- a/src/Extension/Engine/SiteTreePublishingEngine.php
+++ b/src/Extension/Engine/SiteTreePublishingEngine.php
@@ -3,8 +3,8 @@
 namespace SilverStripe\StaticPublishQueue\Extension\Engine;
 
 use SilverStripe\CMS\Model\SiteTree;
-use SilverStripe\CMS\Model\SiteTreeExtension;
 use SilverStripe\Core\Environment;
+use SilverStripe\Core\Extension;
 use SilverStripe\Core\Injector\Injector;
 use SilverStripe\Core\Resettable;
 use SilverStripe\Dev\Deprecation;
@@ -26,9 +26,9 @@
  *
  * @see PublishableSiteTree
  *
- * @extends DataExtension<SiteTree>
+ * @extends Extension<SiteTree>
  */
-class SiteTreePublishingEngine extends SiteTreeExtension implements Resettable
+class SiteTreePublishingEngine extends Extension implements Resettable
 {
     public const ACTION_PUBLISH = 'publish';
     public const ACTION_UNPUBLISH = 'unpublish';
diff --git a/src/Extension/Publishable/PublishableSiteTree.php b/src/Extension/Publishable/PublishableSiteTree.php
index 394cd1f..8ccf6cb 100644
--- a/src/Extension/Publishable/PublishableSiteTree.php
+++ b/src/Extension/Publishable/PublishableSiteTree.php
@@ -6,7 +6,7 @@
 use SilverStripe\CMS\Model\SiteTree;
 use SilverStripe\CMS\Model\VirtualPage;
 use SilverStripe\Control\Director;
-use SilverStripe\ORM\DataExtension;
+use SilverStripe\Core\Extension;
 use SilverStripe\ORM\SS_List;
 use SilverStripe\StaticPublishQueue\Contract\StaticallyPublishable;
 use SilverStripe\StaticPublishQueue\Contract\StaticPublishingTrigger;
@@ -20,9 +20,9 @@
  *
  * @see SiteTreePublishingEngine
  *
- * @extends DataExtension<SiteTree>
+ * @extends Extension<SiteTree>
  */
-class PublishableSiteTree extends DataExtension implements StaticallyPublishable, StaticPublishingTrigger
+class PublishableSiteTree extends Extension implements StaticallyPublishable, StaticPublishingTrigger
 {
     public const REGENERATE_RELATIONS_NONE = 'none';
     public const REGENERATE_RELATIONS_DIRECT = 'direct';