Merge "Provision on different drivers dynamically"

3 files changed, 113 insertions, 2 deletions

diff --git a/lib/sqlalchemy/testing/exclusions.py b/lib/sqlalchemy/testing/exclusions.py
index 1a23ebf41..6ec438193 100644
--- a/lib/sqlalchemy/testing/exclusions.py
+++ b/lib/sqlalchemy/testing/exclusions.py
@@ -272,6 +272,9 @@ class SpecPredicate(Predicate):
     }
 
     def __call__(self, config):
+        if config is None:
+            return False
+
         engine = config.db
 
         if "+" in self.db:
diff --git a/lib/sqlalchemy/testing/plugin/plugin_base.py b/lib/sqlalchemy/testing/plugin/plugin_base.py
index 49ff0f975..b5f2a3e0b 100644
--- a/lib/sqlalchemy/testing/plugin/plugin_base.py
+++ b/lib/sqlalchemy/testing/plugin/plugin_base.py
@@ -17,10 +17,14 @@ is pytest.
 from __future__ import absolute_import
 
 import abc
+import logging
 import re
 import sys
 
 
+log = logging.getLogger("sqlalchemy.testing.plugin_base")
+
+
 py3k = sys.version_info >= (3, 0)
 
 if py3k:
@@ -91,6 +95,15 @@ def setup_options(make_option):
         help="Database uri.  Multiple OK, " "first one is run by default.",
     )
     make_option(
+        "--dbdriver",
+        action="append",
+        type="string",
+        dest="dbdriver",
+        help="Additional database drivers to include in tests.  "
+        "These are linked to the existing database URLs by the "
+        "provisioning system.",
+    )
+    make_option(
         "--dropfirst",
         action="store_true",
         dest="dropfirst",
@@ -350,6 +363,7 @@ def _init_symbols(options, file_config):
 
 @post
 def _engine_uri(options, file_config):
+
     from sqlalchemy.testing import config
     from sqlalchemy import testing
     from sqlalchemy.testing import provision
@@ -359,6 +373,8 @@ def _engine_uri(options, file_config):
     else:
         db_urls = []
 
+    extra_drivers = options.dbdriver or []
+
     if options.db:
         for db_token in options.db:
             for db in re.split(r"[,\s]+", db_token):
@@ -374,7 +390,11 @@ def _engine_uri(options, file_config):
         db_urls.append(file_config.get("db", "default"))
 
     config._current = None
-    for db_url in db_urls:
+
+    expanded_urls = list(provision.generate_db_urls(db_urls, extra_drivers))
+
+    for db_url in expanded_urls:
+        log.info("Adding database URL: %s", db_url)
 
         if options.write_idents and provision.FOLLOWER_IDENT:  # != 'master':
             with open(options.write_idents, "a") as file_:
@@ -596,7 +616,8 @@ def stop_test_class(cls):
 
 
 def _restore_engine():
-    config._current.reset(testing)
+    if config._current:
+        config._current.reset(testing)
 
 
 def final_process_cleanup():
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()