Improved Doctrine setup

Author: ankane

README.md

@@ -117,18 +117,12 @@ Install the package
 composer require pgvector/pgvector
 ```
 
-Register the distance functions
+Register the types and distance functions
 
 ```php
 use Pgvector\Doctrine\PgvectorSetup;
 
-PgvectorSetup::registerFunctions($config);
-```
-
-And the types
-
-```php
-PgvectorSetup::registerTypes($entityManager);
+PgvectorSetup::register($entityManager);
 ```
 
 Enable the extension

src/doctrine/PgvectorSetup.php

@@ -2,29 +2,37 @@
 
 namespace Pgvector\Doctrine;
 
+use Doctrine\DBAL\Platforms\AbstractPlatform;
 use Doctrine\DBAL\Types\Type;
 use Doctrine\ORM\Configuration;
 use Doctrine\ORM\EntityManager;
 
 abstract class PgvectorSetup
 {
-    public static function registerTypes(?EntityManager $entityManager = null): void
+    public static function register(EntityManager $entityManager): void
+    {
+        self::registerTypes();
+        self::registerPlatformTypes($entityManager->getConnection()->getDatabasePlatform());
+        self::registerFunctions($entityManager->getConfiguration());
+    }
+
+    private static function registerTypes(): void
     {
         Type::addType('vector', 'Pgvector\Doctrine\VectorType');
         Type::addType('halfvec', 'Pgvector\Doctrine\HalfVectorType');
         Type::addType('bit', 'Pgvector\Doctrine\BitType');
         Type::addType('sparsevec', 'Pgvector\Doctrine\SparseVectorType');
+    }
 
-        if (!is_null($entityManager)) {
-            $platform = $entityManager->getConnection()->getDatabasePlatform();
-            $platform->registerDoctrineTypeMapping('vector', 'vector');
-            $platform->registerDoctrineTypeMapping('halfvec', 'halfvec');
-            $platform->registerDoctrineTypeMapping('bit', 'bit');
-            $platform->registerDoctrineTypeMapping('sparsevec', 'sparsevec');
-        }
+    private static function registerPlatformTypes(AbstractPlatform $platform): void
+    {
+        $platform->registerDoctrineTypeMapping('vector', 'vector');
+        $platform->registerDoctrineTypeMapping('halfvec', 'halfvec');
+        $platform->registerDoctrineTypeMapping('bit', 'bit');
+        $platform->registerDoctrineTypeMapping('sparsevec', 'sparsevec');
     }
 
-    public static function registerFunctions(Configuration $config): void
+    private static function registerFunctions(Configuration $config): void
     {
         $config->addCustomNumericFunction('l2_distance', 'Pgvector\Doctrine\L2Distance');
         $config->addCustomNumericFunction('max_inner_product', 'Pgvector\Doctrine\MaxInnerProduct');

tests/DoctrineTest.php

@@ -26,7 +26,6 @@ public static function setUpBeforeClass(): void
             paths: [__DIR__ . '/models'],
             isDevMode: true
         );
-        PgvectorSetup::registerFunctions($config);
 
         $connection = DriverManager::getConnection([
             'driver' => 'pgsql',
@@ -35,7 +34,7 @@ public static function setUpBeforeClass(): void
 
         $entityManager = new EntityManager($connection, $config);
         $entityManager->getConnection()->executeStatement('CREATE EXTENSION IF NOT EXISTS vector');
-        PgvectorSetup::registerTypes($entityManager);
+        PgvectorSetup::register($entityManager);
 
         $schemaManager = $entityManager->getConnection()->createSchemaManager();
         try {