diff --git a/src/Extension/MemberExtension.php b/src/Extension/MemberExtension.php
index 5cf29a6..adde78e 100644
--- a/src/Extension/MemberExtension.php
+++ b/src/Extension/MemberExtension.php
@@ -2,13 +2,13 @@
 
 namespace SilverStripe\RealMe\Extension;
 
-use SilverStripe\ORM\DataExtension;
+use SilverStripe\Core\Extension;
 use SilverStripe\Security\Member;
 
 /**
- * @extends DataExtension<Member>
+ * @extends Extension<Member>
  */
-class MemberExtension extends DataExtension
+class MemberExtension extends Extension
 {
     private static $db = array(
         "RealmeSPNameID" => "Varchar(35)",
diff --git a/src/Extension/SiteTreeExtension.php b/src/Extension/SiteTreeExtension.php
index 42b451a..bfb2c00 100644
--- a/src/Extension/SiteTreeExtension.php
+++ b/src/Extension/SiteTreeExtension.php
@@ -7,12 +7,12 @@
 use SilverStripe\Security\InheritedPermissions;
 use SilverStripe\Security\Member;
 use SilverStripe\ORM\DataObject;
-use SilverStripe\ORM\DataExtension;
+use SilverStripe\Core\Extension;
 
 /**
- * @extends DataExtension<SiteTree>
+ * @extends Extension<SiteTree>
  */
-class SiteTreeExtension extends DataExtension
+class SiteTreeExtension extends Extension
 {
     private static $dependencies = array(
         'service' => '%$' . RealMeService::class