backends: declare backends via importlib.metadata entry_points

3 files changed, 24 insertions, 29 deletions

diff --git a/setup.cfg b/setup.cfg
index f314b45..12f68ec 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -28,6 +28,7 @@ packages =
 install_requires = 
     sqlparse
     tabulate
+    importlib_metadata
 
 [options.package_data]
 yoyo = tests/migrations/*.py
@@ -41,3 +42,15 @@ pyodbc = pyodbc
 console_scripts = 
     yoyo = yoyo.scripts.main:main
     yoyo-migrate = yoyo.scripts.main:main
+
+yoyo.backends =
+    odbc = yoyo.backends:ODBCBackend
+    oracle = yoyo.backends:OracleBackend
+    postgresql = yoyo.backends:PostgresqlBackend
+    postgres = yoyo.backends:PostgresqlBackend
+    psql = yoyo.backends:PostgresqlBackend
+    mysql = yoyo.backends:MySQLBackend
+    mysql+mysqldb = yoyo.backends:MySQLdbBackend
+    sqlite = yoyo.backends:SQLiteBackend
+    snowflake = yoyo.backends:SnowflakeBackend
+    redshift = yoyo.backends:RedshiftBackend
diff --git a/yoyo/backends.py b/yoyo/backends.py
index c6b6f21..dc60130 100644
--- a/yoyo/backends.py
+++ b/yoyo/backends.py
@@ -19,6 +19,7 @@ from importlib import import_module
 from itertools import count
 from logging import getLogger
 from typing import Dict
+from importlib_metadata import entry_points
 
 import getpass
 import os
@@ -791,6 +792,11 @@ class SnowflakeBackend(DatabaseBackend):
         pass
 
 
+def get_backend_class(name):
+    backend_eps = entry_points(group="yoyo.backends")
+    return backend_eps[name].load()
+
+
 def get_dbapi_module(name):
     """
     Import and return the named DB-API driver module
diff --git a/yoyo/connections.py b/yoyo/connections.py
index 3df118a..24ea5ea 100755
--- a/yoyo/connections.py
+++ b/yoyo/connections.py
@@ -20,28 +20,8 @@ from urllib.parse import urlencode
 from urllib.parse import urlsplit
 from urllib.parse import urlunsplit
 
-from .migrations import default_migration_table
-from .backends import PostgresqlBackend
-from .backends import RedshiftBackend
-from .backends import SQLiteBackend
-from .backends import ODBCBackend
-from .backends import OracleBackend
-from .backends import MySQLBackend
-from .backends import MySQLdbBackend
-from .backends import SnowflakeBackend
-
-BACKENDS = {
-    "odbc": ODBCBackend,
-    "oracle": OracleBackend,
-    "postgresql": PostgresqlBackend,
-    "postgres": PostgresqlBackend,
-    "psql": PostgresqlBackend,
-    "mysql": MySQLBackend,
-    "mysql+mysqldb": MySQLdbBackend,
-    "sqlite": SQLiteBackend,
-    "snowflake": SnowflakeBackend,
-    "redshift": RedshiftBackend,
-}
+from yoyo.backends import get_backend_class
+from yoyo.migrations import default_migration_table
 
 
 _DatabaseURI = namedtuple(
@@ -91,7 +71,7 @@ def get_backend(uri, migration_table=default_migration_table):
     """
     parsed = parse_uri(uri)
     try:
-        backend_class = BACKENDS[parsed.scheme.lower()]
+        backend_class = get_backend_class(parsed.scheme.lower())
     except KeyError:
         raise BadConnectionURI(
             "Unrecognised database connection scheme %r" % parsed.scheme
@@ -119,12 +99,8 @@ def parse_uri(s):
 
     return DatabaseURI(
         scheme=result.scheme,
-        username=(
-            unquote(result.username) if result.username is not None else None
-        ),
-        password=(
-            unquote(result.password) if result.password is not None else None
-        ),
+        username=(unquote(result.username) if result.username is not None else None),
+        password=(unquote(result.password) if result.password is not None else None),
         hostname=result.hostname,
         port=result.port,
         database=result.path[1:] if result.path else None,