build: make the database drivers optional extras

CONTRIBUTING.md

@@ -54,12 +54,13 @@ pixi run -e test pytest -m postgresql            # PostgreSQL tests only
 export DOCKER_HOST=unix://$HOME/.docker/run/docker.sock
 ```
 
-### PostgreSQL Backend
+### Database Backends
 
-DataJoint supports MySQL 8.0.13+ and PostgreSQL 15+ as production database backends. To install the PostgreSQL driver:
+DataJoint supports PostgreSQL 15+ and MySQL 8.0.13+ as production database backends. Each database driver is an optional extra, and the `test` and `dev` extras install both:
 
 ```bash
 pip install -e ".[postgres]"    # Installs psycopg2-binary
+pip install -e ".[mysql]"       # Installs pymysql
 ```
 
 Tests automatically spin up both MySQL and PostgreSQL containers via testcontainers. Backend-parameterized tests (those using the `backend` fixture in `tests/conftest.py`) run against both backends to ensure feature parity.

Dockerfile

@@ -19,6 +19,6 @@ USER root
 RUN \
     chown -R ${HOST_UID:-1000}:mambauser /main && \
     chown -R ${HOST_UID:-1000}:mambauser /src && \
-    ${CONDA_BIN} run -n base pip install -q --no-cache-dir /main && \
+    ${CONDA_BIN} run -n base pip install -q --no-cache-dir "/main[postgres,mysql]" && \
     rm -r /main/*
 USER ${MAMBA_USER}

README.md

@@ -57,8 +57,11 @@ The **DataJoint Library** is a framework for scientific data pipelines based on
 
 ## Installation
 
+Install DataJoint with the driver for your database:
+
 ```bash
-pip install datajoint
+pip install "datajoint[postgres]"   # PostgreSQL
+pip install "datajoint[mysql]"      # MySQL
 ```
 
 or with Conda:

RELEASE_MEMO.md

@@ -143,7 +143,9 @@ If the bot doesn't create a PR, manually update the feedstock:
    curl -sL https://pypi.org/pypi/datajoint/2.1.0/json | jq -r '.urls[] | select(.packagetype=="sdist") | .digests.sha256'
    ```
 
-4. **Check dependencies** match `pyproject.toml`:
+4. **Check dependencies** match `pyproject.toml`. conda has no extras, so `run`
+   also lists the database drivers that `pyproject.toml` keeps in the `mysql` and
+   `postgres` extras:
    ```yaml
    requirements:
      host:
@@ -154,6 +156,7 @@ If the bot doesn't create a PR, manually update the feedstock:
        - python >={{ python_min }}
        - numpy
        - pymysql >=0.7.2
+       - psycopg2 >=2.9.0
        - deepdiff
        - pyparsing
        - pandas

pyproject.toml

@@ -7,7 +7,6 @@ name = "datajoint"
 dynamic = ["version"]
 dependencies = [
   "numpy",
-  "pymysql>=0.7.2",
   "deepdiff",
   "pyparsing",
   "pandas",
@@ -93,6 +92,8 @@ test = [
   "ipython",
   "graphviz",
   "testcontainers[mysql,minio,postgres]>=4.0",
+  "pymysql>=0.7.2",
+  "psycopg2-binary>=2.9.0",
   "polars>=0.20.0",
   "pyarrow>=14.0.0",
 ]
@@ -101,7 +102,9 @@ test = [
 s3 = ["s3fs>=2023.1.0"]
 gcs = ["gcsfs>=2023.1.0"]
 azure = ["adlfs>=2023.1.0"]
+# Database drivers: install the extra for each backend in use.
 postgres = ["psycopg2-binary>=2.9.0"]
+mysql = ["pymysql>=0.7.2"]
 polars = ["polars>=0.20.0"]
 arrow = ["pyarrow>=14.0.0"]
 viz = ["matplotlib", "ipython"]
@@ -114,6 +117,7 @@ test = [
   "ipython",
   "s3fs>=2023.1.0",
   "testcontainers[mysql,minio,postgres]>=4.0",
+  "pymysql>=0.7.2",
   "psycopg2-binary>=2.9.0",
   "polars>=0.20.0",
   "pyarrow>=14.0.0",
@@ -125,6 +129,8 @@ dev = [
   # including test
   "pytest",
   "pytest-cov",
+  "pymysql>=0.7.2",
+  "psycopg2-binary>=2.9.0",
   "polars>=0.20.0",
   "pyarrow>=14.0.0",
 ]

src/datajoint/__init__.py

@@ -278,7 +278,7 @@ def FreeTable(conn_or_name, full_table_name: str | None = None) -> _FreeTable:
 # =============================================================================
 # Lazy imports — heavy dependencies loaded on first access
 # =============================================================================
-# These modules import heavy dependencies (networkx, matplotlib, click, pymysql)
+# These modules import heavy dependencies (networkx, matplotlib, click)
 # that slow down `import datajoint`. They are loaded on demand.
 
 _lazy_modules = {

src/datajoint/adapters/base.py

@@ -7,9 +7,48 @@
 
 from __future__ import annotations
 
+import importlib
 from abc import ABC, abstractmethod
+from types import ModuleType
 from typing import Any
 
+from ..errors import DataJointError
+
+
+def import_driver(module: str, backend: str, extra: str) -> ModuleType:
+    """
+    Import a database driver, or name the extra that installs it.
+
+    Database drivers are optional extras, so adapters import them when they are
+    created rather than when ``datajoint`` is imported.
+
+    Parameters
+    ----------
+    module : str
+        Import name of the driver, e.g. ``'pymysql'`` or ``'psycopg2'``.
+    backend : str
+        Backend name for the error message, e.g. ``'MySQL'``.
+    extra : str
+        The ``datajoint`` extra that installs the driver, e.g. ``'mysql'``.
+
+    Returns
+    -------
+    ModuleType
+        The imported driver module.
+
+    Raises
+    ------
+    DataJointError
+        If the driver is not installed.
+    """
+    try:
+        return importlib.import_module(module)
+    except ImportError as error:
+        raise DataJointError(
+            f"The {backend} backend requires the {module} driver, which is not installed. "
+            f'Install it with: pip install "datajoint[{extra}]"'
+        ) from error
+
 
 class DatabaseAdapter(ABC):
     """

src/datajoint/adapters/mysql.py

@@ -9,10 +9,12 @@
 
 from typing import Any
 
-import pymysql as client
-
 from .. import errors
-from .base import DatabaseAdapter
+from .base import DatabaseAdapter, import_driver
+
+# The pymysql driver (the ``datajoint[mysql]`` extra), imported when the first
+# MySQLAdapter is created so that ``import datajoint`` loads no database driver.
+client: Any = None
 
 # Core type mapping: DataJoint core types → MySQL types
 CORE_TYPE_MAP = {
@@ -48,6 +50,18 @@
 class MySQLAdapter(DatabaseAdapter):
     """MySQL database adapter implementation."""
 
+    def __init__(self) -> None:
+        """
+        Initialize MySQL adapter.
+
+        Raises
+        ------
+        DataJointError
+            If pymysql is not installed (``pip install "datajoint[mysql]"``).
+        """
+        global client
+        client = import_driver("pymysql", backend="MySQL", extra="mysql")
+
     # =========================================================================
     # Connection Management
     # =========================================================================
@@ -175,9 +189,7 @@ def get_cursor(self, connection: Any, as_dict: bool = False) -> Any:
         Any
             pymysql cursor object.
         """
-        import pymysql
-
-        cursor_class = pymysql.cursors.DictCursor if as_dict else pymysql.cursors.Cursor
+        cursor_class = client.cursors.DictCursor if as_dict else client.cursors.Cursor
         return connection.cursor(cursor=cursor_class)
 
     # =========================================================================

src/datajoint/adapters/postgres.py

@@ -10,15 +10,13 @@
 import re
 from typing import Any
 
-try:
-    import psycopg2 as client
-    from psycopg2 import sql
-except ImportError:
-    client = None  # type: ignore
-    sql = None  # type: ignore
-
 from .. import errors
-from .base import DatabaseAdapter
+from .base import DatabaseAdapter, import_driver
+
+# The psycopg2 driver (the ``datajoint[postgres]`` extra), imported when the first
+# PostgreSQLAdapter is created so that ``import datajoint`` loads no database driver.
+# Either the psycopg2 or the psycopg2-binary distribution provides it.
+client: Any = None
 
 # Core type mapping: DataJoint core types → PostgreSQL types
 CORE_TYPE_MAP = {
@@ -56,11 +54,16 @@ class PostgreSQLAdapter(DatabaseAdapter):
     """PostgreSQL database adapter implementation."""
 
     def __init__(self) -> None:
-        """Initialize PostgreSQL adapter."""
-        if client is None:
-            raise ImportError(
-                "psycopg2 is required for PostgreSQL support. " "Install it with: pip install 'datajoint[postgres]'"
-            )
+        """
+        Initialize PostgreSQL adapter.
+
+        Raises
+        ------
+        DataJointError
+            If psycopg2 is not installed (``pip install "datajoint[postgres]"``).
+        """
+        global client
+        client = import_driver("psycopg2", backend="PostgreSQL", extra="postgres")
 
     # =========================================================================
     # Connection Management

tests/unit/test_adapters.py

@@ -4,9 +4,14 @@
 Tests adapter functionality without requiring actual database connections.
 """
 
+import re
+import sys
+
 import pytest
 
 from datajoint.adapters import DatabaseAdapter, MySQLAdapter, PostgreSQLAdapter, get_adapter
+from datajoint.connection import Connection
+from datajoint.errors import DataJointError
 
 
 class TestAdapterRegistry:
@@ -48,6 +53,45 @@ def test_get_adapter_invalid(self):
             get_adapter("sqlite")
 
 
+# Each backend, the driver module its adapter imports, and the extra that installs it.
+MISSING_DRIVER_CASES = [
+    ("mysql", "pymysql", "mysql"),
+    ("postgresql", "psycopg2", "postgres"),
+    ("postgres", "psycopg2", "postgres"),
+]
+
+
+class TestMissingDriver:
+    """A missing database driver is reported with the extra that installs it."""
+
+    @pytest.mark.parametrize("backend, driver, extra", MISSING_DRIVER_CASES)
+    def test_get_adapter_names_extra(self, monkeypatch, backend, driver, extra):
+        """Creating the adapter without its driver raises a DataJointError naming the extra."""
+        # A None entry in sys.modules makes `import <driver>` raise ImportError.
+        monkeypatch.setitem(sys.modules, driver, None)
+        with pytest.raises(DataJointError, match=re.escape(f'pip install "datajoint[{extra}]"')) as info:
+            get_adapter(backend)
+        assert driver in str(info.value)
+        assert isinstance(info.value.__cause__, ImportError)
+
+    @pytest.mark.parametrize("backend, driver, extra", MISSING_DRIVER_CASES)
+    def test_connection_names_extra(self, monkeypatch, backend, driver, extra):
+        """Connecting without the driver fails before contacting any server."""
+        monkeypatch.setitem(sys.modules, driver, None)
+        with pytest.raises(DataJointError, match=re.escape(f'pip install "datajoint[{extra}]"')):
+            Connection("localhost", "user", "password", backend=backend)
+
+    def test_other_driver_not_required(self, monkeypatch):
+        """Each backend needs only its own driver."""
+        pytest.importorskip("psycopg2")
+        with monkeypatch.context() as patch:
+            patch.setitem(sys.modules, "pymysql", None)
+            assert isinstance(get_adapter("postgresql"), PostgreSQLAdapter)
+        with monkeypatch.context() as patch:
+            patch.setitem(sys.modules, "psycopg2", None)
+            assert isinstance(get_adapter("mysql"), MySQLAdapter)
+
+
 class TestMySQLAdapter:
     """Test MySQL adapter implementation."""
 

tests/unit/test_lazy_imports.py

@@ -2,9 +2,11 @@
 Tests for lazy import behavior.
 
 These tests verify that heavy dependencies (networkx, matplotlib, click)
-are not loaded until their associated features are accessed.
+are not loaded until their associated features are accessed, and that
+``import datajoint`` loads no database driver.
 """
 
+import subprocess
 import sys
 
 
@@ -85,3 +87,20 @@ def test_core_imports_available():
     # Heavy modules should still not be loaded
     assert "datajoint.diagram" not in sys.modules
     assert "datajoint.cli" not in sys.modules
+
+
+def test_import_loads_no_database_driver():
+    """import datajoint should not import pymysql or psycopg2.
+
+    Drivers are optional extras, so the import must succeed without them and
+    leave the driver to the adapter that connects. A fresh interpreter is used
+    so that drivers already imported by other tests do not hide the import.
+    """
+    code = (
+        "import sys, warnings\n"
+        "warnings.simplefilter('ignore')\n"
+        "import datajoint\n"
+        "print(sorted(m for m in sys.modules if m.split('.')[0] in ('pymysql', 'psycopg2')))\n"
+    )
+    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True)
+    assert result.stdout.strip() == "[]", f"import datajoint loaded database drivers: {result.stdout.strip()}"