From b8ee7f07896e55a09c1af52e854f4596645e0a4f Mon Sep 17 00:00:00 2001 From: Dimitri Yatsenko Date: Sat, 17 Oct 2026 18:39:24 -0500 Subject: [PATCH] build: make the database drivers optional extras pymysql was a core dependency, so every install of datajoint pulled in the MySQL driver whether or not it was used, while the PostgreSQL driver was already the optional postgres extra. Both drivers are now extras: pip install "datajoint[postgres]" installs psycopg2-binary, and pip install "datajoint[mysql]" installs pymysql, with the same >=0.7.2 floor as before. A bare pip install datajoint installs no database driver. import datajoint no longer imports a driver. Each adapter imports its driver when it is created, which happens when a connection is made, through a new helper, import_driver, in adapters/base.py. When the driver is missing, connecting raises a DataJointError that names the extra, for example: The MySQL backend requires the pymysql driver, which is not installed. Install it with: pip install "datajoint[mysql]". The PostgreSQL adapter raised ImportError for its missing driver; it now raises the same DataJointError. The adapter imports the psycopg2 module, which both the psycopg2 and the psycopg2-binary distributions provide, so either one satisfies it. The test dependency group and the test and dev extras install both drivers. The README install line, CONTRIBUTING.md, the development Docker image (it now installs both extras), and the conda-forge notes in RELEASE_MEMO.md (conda has no extras, so the recipe lists both drivers) are updated to match. Migration: an environment that connects to MySQL needs pip install "datajoint[mysql]", or pymysql installed some other way; one that connects to PostgreSQL needs pip install "datajoint[postgres]". An upgrade in place keeps a pymysql that is already installed, but a fresh install from a requirements file that lists only datajoint gets no driver. Code that caught ImportError from PostgreSQLAdapter() should catch DataJointError instead. Tests: test_lazy_imports.py checks, in a fresh interpreter, that import datajoint loads neither pymysql nor psycopg2. test_adapters.py blocks each driver in sys.modules and checks that get_adapter and Connection raise a DataJointError naming the right extra before any server is contacted, and that each backend needs only its own driver. --- CONTRIBUTING.md | 5 ++-- Dockerfile | 2 +- README.md | 5 +++- RELEASE_MEMO.md | 5 +++- pyproject.toml | 8 +++++- src/datajoint/__init__.py | 2 +- src/datajoint/adapters/base.py | 39 ++++++++++++++++++++++++++ src/datajoint/adapters/mysql.py | 24 ++++++++++++---- src/datajoint/adapters/postgres.py | 29 +++++++++++--------- tests/unit/test_adapters.py | 44 ++++++++++++++++++++++++++++++ tests/unit/test_lazy_imports.py | 21 +++++++++++++- 11 files changed, 157 insertions(+), 27 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 64ad3f6ab..9b725a90d 100644 --- a/CONTRIBUTING.md +++ b/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. diff --git a/Dockerfile b/Dockerfile index 88c046345..28c94f18a 100644 --- a/Dockerfile +++ b/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} diff --git a/README.md b/README.md index 7b68d185b..099864cc3 100644 --- a/README.md +++ b/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: diff --git a/RELEASE_MEMO.md b/RELEASE_MEMO.md index 1a3623330..951efde55 100644 --- a/RELEASE_MEMO.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 050bde152..358299189 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/src/datajoint/__init__.py b/src/datajoint/__init__.py index 69c6df269..e6b0c6f66 100644 --- a/src/datajoint/__init__.py +++ b/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 = { diff --git a/src/datajoint/adapters/base.py b/src/datajoint/adapters/base.py index 5b4a7ab72..0c4e8e94d 100644 --- a/src/datajoint/adapters/base.py +++ b/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): """ diff --git a/src/datajoint/adapters/mysql.py b/src/datajoint/adapters/mysql.py index 2ef9baf84..857913128 100644 --- a/src/datajoint/adapters/mysql.py +++ b/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) # ========================================================================= diff --git a/src/datajoint/adapters/postgres.py b/src/datajoint/adapters/postgres.py index 3a55d8eb8..5e18470ac 100644 --- a/src/datajoint/adapters/postgres.py +++ b/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 diff --git a/tests/unit/test_adapters.py b/tests/unit/test_adapters.py index bd0afa9e0..673977901 100644 --- a/tests/unit/test_adapters.py +++ b/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 ` 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.""" diff --git a/tests/unit/test_lazy_imports.py b/tests/unit/test_lazy_imports.py index a87412151..36e33c68d 100644 --- a/tests/unit/test_lazy_imports.py +++ b/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()}"