1 files changed, 87 insertions, 0 deletions
diff --git a/lib/sqlalchemy/testing/provision.py b/lib/sqlalchemy/testing/provision.py
index 660a3bd24..13a5ea078 100644
--- a/lib/sqlalchemy/testing/provision.py
+++ b/lib/sqlalchemy/testing/provision.py
@@ -1,8 +1,10 @@
 import collections
+import copy
 import logging
 
 from . import config
 from . import engines
+from .. import exc
 from ..engine import url as sa_url
 from ..util import compat
 
@@ -73,6 +75,91 @@ def drop_follower_db(follower_ident):
         drop_db(cfg, cfg.db, follower_ident)
 
 
+def generate_db_urls(db_urls, extra_drivers):
+    """Generate a set of URLs to test given configured URLs plus additional
+    driver names.
+
+    Given::
+
+        --dburi postgresql://db1  \
+        --dburi postgresql://db2  \
+        --dburi postgresql://db2  \
+        --dbdriver=psycopg2 --dbdriver=asyncpg
+
+    Noting that the default postgresql driver is psycopg2.  the output
+    would be::
+
+        postgresql+psycopg2://db1
+        postgresql+asyncpg://db1?async_fallback=true
+        postgresql+psycopg2://db2
+        postgresql+psycopg2://db3
+
+    That is, for the driver in a --dburi, we want to keep that and use that
+    driver for each URL it's part of .   For a driver that is only
+    in --dbdrivers, we want to use it just once for one of the URLs.
+    for a driver that is both coming from --dburi as well as --dbdrivers,
+    we want to keep it in that dburi.
+
+
+    """
+    urls = set()
+
+    backend_to_driver_we_already_have = collections.defaultdict(set)
+
+    urls_plus_dialects = [
+        (url_obj, url_obj.get_dialect())
+        for url_obj in [sa_url.make_url(db_url) for db_url in db_urls]
+    ]
+
+    for url_obj, dialect in urls_plus_dialects:
+        backend_to_driver_we_already_have[dialect.name].add(dialect.driver)
+
+    backend_to_driver_we_need = {}
+
+    for url_obj, dialect in urls_plus_dialects:
+        backend = dialect.name
+        dialect.load_provisioning()
+
+        if backend not in backend_to_driver_we_need:
+            backend_to_driver_we_need[backend] = extra_per_backend = set(
+                extra_drivers
+            ).difference(backend_to_driver_we_already_have[backend])
+        else:
+            extra_per_backend = backend_to_driver_we_need[backend]
+
+        for driver_url in _generate_driver_urls(url_obj, extra_per_backend):
+            if driver_url in urls:
+                continue
+            urls.add(driver_url)
+            yield driver_url
+
+
+def _generate_driver_urls(url, extra_drivers):
+    main_driver = url.get_driver_name()
+    extra_drivers.discard(main_driver)
+
+    yield str(url)
+
+    for drv in list(extra_drivers):
+        new_url = generate_driver_url(url, drv)
+        if new_url:
+            extra_drivers.remove(drv)
+            yield str(new_url)
+
+
+@register.init
+def generate_driver_url(url, driver):
+    backend = url.get_backend_name()
+    new_url = copy.copy(url)
+    new_url.drivername = "%s+%s" % (backend, driver)
+    try:
+        new_url.get_dialect()
+    except exc.NoSuchModuleError:
+        return None
+    else:
+        return new_url
+
+
 def _configs_for_db_operation():
     hosts = set()