Add create_engine function for SQLAlchemy

Author: kesmit13

docs/src/api.rst

@@ -15,10 +15,17 @@ package. It connects to a SingleStore database using either
 `DB-API <https://peps.python.org/pep-0249/>`_ compliant parameters,
 or a connection string in the form of a URL.
 
+The :func:`create_engine` function is used with the SQLAlchemy package
+to create an SQLAlchemy engine for SingleStoreDB connections. This is
+primarily for use in environments where the connection parameters are
+stored in environment variables so that you can create SingleStoreDB
+connections without specifying any parameters in the code itself.
+
 .. autosummary::
    :toctree: generated/
 
    connect
+   create_engine
 
 
 Connection

singlestoredb/__init__.py

@@ -15,6 +15,7 @@
 
 __version__ = '0.8.0'
 
+from .alchemy import create_engine
 from .config import options, get_option, set_option, describe_option
 from .connection import connect, apilevel, threadsafety, paramstyle
 from .exceptions import (

singlestoredb/alchemy/__init__.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+import inspect
+from typing import Any
+from urllib.parse import quote
+
+try:
+    import sqlalchemy_singlestoredb  # noqa: F401
+    import sqlalchemy
+    has_sqlalchemy = True
+except ImportError:
+    has_sqlalchemy = False
+
+from ..connection import build_params
+from ..connection import connect
+
+
+def create_engine(*args: Any, **kwargs: Any) -> Any:
+    """
+    Create an SQLAlchemy engine for SingleStoreDB.
+
+    Parameters
+    ----------
+    **kwargs : Any
+        The parameters taken here are the same as for
+        `sqlalchemy.create_engine`. However, this function can be
+        called without any parameters in order to inherit parameters
+        set by environment variables or parameters set in by
+        options in Python code.
+
+    See Also
+    --------
+    `sqlalchemy.create_engine`
+
+    Returns
+    -------
+    SQLAlchemy engine
+
+    """
+    if not has_sqlalchemy:
+        raise RuntimeError('sqlalchemy_singlestoredb package is not installed')
+
+    if len(args) > 1:
+        raise ValueError(
+            '`args` can only have a single element '
+            'containing the database URL',
+        )
+
+    if args:
+        kwargs['host'] = args[0]
+
+    conn_params = {}
+    sa_params = {}
+
+    conn_args = inspect.getfullargspec(connect).args
+
+    for key, value in kwargs.items():
+        if key in conn_args:
+            conn_params[key] = value
+        else:
+            sa_params[key] = value
+
+    params = build_params(**conn_params)
+    driver = params.pop('driver', None)
+    host = params.pop('host')
+    port = params.pop('port')
+    user = params.pop('user', None)
+    password = params.pop('password', None)
+    database = params.pop('database', '')
+
+    if not driver:
+        driver = 'singlestoredb+mysql'
+    elif not driver.startswith('singlestoredb'):
+        driver = f'singlestoredb+{driver}'
+
+    if user is not None and password is not None:
+        url = f'{driver}://{quote(user)}:{quote(password)}@' \
+              f'{host}:{port}/{quote(database)}'
+    elif user is not None:
+        url = f'{driver}://{quote(user)}@{host}:{port}/{quote(database)}'
+    elif password is not None:
+        url = f'{driver}://:{quote(password)}@{host}:{port}/{quote(database)}'
+    else:
+        url = f'{driver}://{host}:{port}/{quote(database)}'
+
+    return sqlalchemy.create_engine(url, connect_args=params, **sa_params)