import os
from contextlib import suppress
from importlib import import_module
import sqlalchemy as sa
from django.conf import settings
from django.utils.encoding import force_str
from django.utils.module_loading import import_string
DIALECT_MAP = {
"django.db.backends.sqlite3": "sqlite",
"django.db.backends.postgresql": "postgresql",
"django.db.backends.mysql": "mysql",
"django.db.backends.oracle": "oracle",
"sqlserver_ado": "mssql",
}
DIALECT_MAP_TO_DJANGO = {v: k for k, v in DIALECT_MAP.items()}
[docs]def boolean(x):
return str(x) in {"True", "1"}
[docs]def integer(x):
return int(x)
[docs]def string(x):
return force_str(x)
[docs]def string_list(x):
return force_str(x).split(",")
[docs]def importable(x):
try:
return import_string(x)
except ImportError:
return import_module(x)
[docs]def importable_list(x):
return [importable(i) for i in force_str(x).split(",")]
[docs]def importable_list_tuples(x):
return [(importable(i), j) for i, j in [k.split(":") for k in force_str(x).split(",")]]
ENGINE_OPTIONS_NORMALIZATION = {
"case_sensitive": boolean,
# connect_args - direct querystring params in url
"convert_unicode": boolean,
"creator": importable,
"echo": boolean,
"echo_pool": boolean,
"empty_in_strategy": string,
"encoding": string,
# execution_options - signal should be used instead
"implicit_returning": boolean,
"isolation_level": string,
"label_length": integer,
"listeners": importable_list,
"logging_name": string,
"max_overflow": integer,
"module": importable,
"paramstyle": string,
"plugins": string_list,
"pool": lambda x: importable(x)(),
"pool_events": importable_list_tuples,
"pool_logging_name": string,
"pool_pre_ping": boolean,
"pool_recycle": integer,
"pool_reset_on_return": string,
"pool_size": integer,
"pool_threadlocal": boolean,
"pool_timeout": integer,
"pool_use_lifo": boolean,
"poolclass": importable,
"strategy": string,
}
[docs]def get_settings(alias):
"""Returns database settings from either ``SQLALCHEMY_CONNECTIONS`` setting
or ``DATABASES`` setting."""
if hasattr(settings, "SQLALCHEMY_CONNECTIONS") and alias in settings.SQLALCHEMY_CONNECTIONS:
return settings.SQLALCHEMY_CONNECTIONS[alias]
return settings.DATABASES[alias]
def _asdict(url):
return {
"drivername": url.drivername,
"username": url.username,
"password": getattr(url, "password_original", None) or getattr(url, "password", None),
"host": url.host,
"port": url.port,
"database": url.database,
"query": dict(url.query.items()),
}
def _options_from_url(url, base_options):
options = base_options.copy()
url_kwargs = _asdict(url)
query_options = url_kwargs["query"]
options.update(
{
"engine_options": {
k.replace("engine_", ""): ENGINE_OPTIONS_NORMALIZATION.get(k.replace("engine_", ""), lambda x: x)(
query_options.pop(k)
)
for k in list(query_options)
if k.startswith("engine_")
}
}
)
return sa.engine.url.URL(**url_kwargs), options
[docs]def make_url(alias_or_url):
"""
Generates a URL either from a url string, environment variable, SQLALCHEMY_CONNECTIONS or DATABASES settings
---------------------------------------------------------------------
alias_or_url: str
name of the alias or url as string
"""
settings_kwargs = {}
with suppress(KeyError):
_, settings_kwargs = make_url_from_settings(alias_or_url)
try:
url = sa.engine.url.make_url(alias_or_url)
except sa.exc.ArgumentError:
pass
else:
return _options_from_url(url, settings_kwargs)
alias = alias_or_url
url = sa.engine.url.make_url(os.environ.get(f"{alias.upper()}_URL", None))
if url:
return _options_from_url(url, settings_kwargs)
return make_url_from_settings(alias)
[docs]def make_url_from_settings(alias):
"""
Generates a URL using the alias in settings
-------------------------------------------
alias: str
name of the alias
Overall settings are very similar with django database settings with a few extra keys.
``USER`` - database user
``PASSWORD`` - database user password
``HOST`` - database host
``NAME`` - database name
``PORT`` - database name
``DIALECT`` - dialect to be used in url, if not provided, will use the ``DIALECT_MAP`` to figure out a dialect to
be used in sqlalchemy url
``DRIVER`` - If provided, will be used as the driver in sqlalchemy url
``SQLALCHEMY`` - If provided, a custom :py:class:`..sqlalchemy.SQLAlchemy` class to be used
``QUERY`` - querystring arguments for sqlalchemy url
``ALCHEMY_OPTIONS`` - Optional arguments to be used to initialize the :py:class:`..sqlalchemy.SQLAlchemy` instance
* ``session_class`` - a custom session class to be used
* ``registry_class`` - a custom registy class to be used for scoping
* ``model_class`` - a custom base model class to be used for declarative base models.
* ``metadata_class`` - a custom metadata class used in delclarative models.
* ``metadata_options`` - custom options to use in metadata creation such as specifying naming conventions.
* ``engine_options`` - arguments for sqlalchemy ``create_engine``
* ``session_options`` - arguments for sqlalchemy ``sessionmaker``
Other options are ignored.
"""
data = get_settings(alias)
if "DIALECT" not in data:
data["DIALECT"] = DIALECT_MAP.get(data["ENGINE"]) or data["ENGINE"].split(".")[-1]
names = [data["DIALECT"].lower()]
if "DRIVER" in data:
names.append(data["DRIVER"].lower())
drivername = "+".join(names)
kwargs = {
"username": data.get("USER") or None,
"password": data.get("PASSWORD") or None,
"host": data.get("HOST") or None,
"database": data.get("NAME") or None,
"query": data.get("QUERY", {}),
}
with suppress(Exception):
kwargs["port"] = int(data.get("PORT"))
url = sa.engine.url.URL(drivername, **kwargs)
return url, data.get("ALCHEMY_OPTIONS", {})