db optional dependency

.github/workflows/dev.yml

@@ -35,7 +35,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
+          python -m pip install -e '.[db]' --upgrade pip
           pip install poetry tox tox-gh-actions
 
       - name: test with tox

Dockerfile

@@ -55,7 +55,7 @@ ENV POETRY_INSTALLER_MAX_WORKERS=10
 
 # Install project in editable mode and with development dependencies
 WORKDIR $APP_PATH
-RUN poetry install
+RUN poetry install --all-extras
 
 ENTRYPOINT ["poetry", "run"]
 CMD ["nml"]

nannyml/__init__.py

@@ -39,7 +39,7 @@
 
 from .calibration import Calibrator, IsotonicCalibrator, needs_calibration
 from .chunk import Chunk, Chunker, CountBasedChunker, DefaultChunker, PeriodBasedChunker, SizeBasedChunker
-from .data_quality import MissingValuesCalculator, UnseenValuesCalculator, NumericalRangeCalculator
+from .data_quality import MissingValuesCalculator, NumericalRangeCalculator, UnseenValuesCalculator
 from .datasets import (
     load_modified_california_housing_dataset,
     load_synthetic_binary_classification_dataset,
@@ -59,7 +59,14 @@
     UnivariateDriftCalculator,
 )
 from .exceptions import ChunkerException, InvalidArgumentsException, MissingMetadataException
-from .io import DatabaseWriter, PickleFileWriter, RawFilesWriter
+from .io import PickleFileWriter, RawFilesWriter
+
+try:
+    from .io.db import DatabaseWriter
+except ImportError:
+    logging.getLogger().warning(
+        "`db` module is not available. Install the `nannyml[db]` extra to use this functionality."
+    )
 from .performance_calculation import PerformanceCalculator
 from .performance_estimation import CBPE, DLE
 from .stats import (

nannyml/io/__init__.py

@@ -14,11 +14,16 @@
 """
 
 from .base import Reader, Writer, WriterFactory
-from .db import DatabaseWriter
 from .file_reader import FileReader
 from .file_writer import FileWriter
 from .pickle_file_writer import PickleFileWriter
 from .raw_files_writer import RawFilesWriter
 from .store import FilesystemStore, JoblibPickleSerializer, Serializer, Store
 
+try:
+    from .db import DatabaseWriter
+except ImportError:
+    print("`db` module is not available. Install the `nannyml[db]` extra to use this functionality.")
+
+
 DEFAULT_WRITER = RawFilesWriter(path='out')

nannyml/io/db/database_writer.py

@@ -1,6 +1,11 @@
 from typing import Any, Dict, Optional
 
-from sqlmodel import Session, SQLModel, create_engine, select
+try:
+    from sqlmodel import Session, SQLModel, create_engine, select
+except ImportError:
+    raise ImportError(
+        "`sqlmodel` module is not available. Please install the `nannyml[db]` extra to use this functionality."
+    )
 
 from nannyml._typing import Result
 from nannyml.exceptions import WriterException

nannyml/io/db/entities.py

@@ -9,12 +9,17 @@
 """
 
 import sys
-
 from datetime import datetime
 from typing import List, Optional
 
 from pydantic import ConfigDict
-from sqlmodel import Field, Relationship, SQLModel
+
+try:
+    from sqlmodel import Field, Relationship, SQLModel
+except ImportError:
+    raise ImportError(
+        "`sqlmodel` module is not available. Please install the `nannyml[db]` extra to use this functionality."
+    )
 
 
 class Model(SQLModel, table=True):  # type: ignore[call-arg]

pyproject.toml

@@ -1,7 +1,7 @@
 [tool]
 [tool.poetry]
 name = "nannyml"
-version = "0.11.0"
+version = "0.12.0"
 homepage = "https://github.com/nannyml/nannyml"
 description = "NannyML, Your library for monitoring model performance."
 authors = ["Niels Nuyttens <niels@nannyml.com>"]
@@ -43,7 +43,7 @@ numpy = "^1.24,<1.25"
 seaborn = "^0.13.2"
 kaleido = "0.2.1"
 pyarrow = "^14.0.0"
-gcsfs = "^2022.5.0"
+gcsfs = ">=2022.5.0"
 pydantic = "^2.7.4"
 rich = "^12.5.1"
 click = "^8.1.3"
@@ -52,15 +52,18 @@ Jinja2 = "<3.1"
 pyfiglet = "^0.8.post1"
 lightgbm = "^3.3.2"
 FLAML = "^1.0.11"
-s3fs = "^2022.8.2"
-sqlmodel = "^0.0.19"
+s3fs = ">=2022.8.2"
+sqlmodel = {version = "^0.0.19", optional = true}
 APScheduler = "^3.9.1"
-psycopg2-binary = "^2.9.3"
+psycopg2-binary = {version = "^2.9.3", optional = true}
 segment-analytics-python = "^2.3.2"
 python-dotenv = "^0.21.0"
 types-pyyaml = "^6.0.12.8"
 types-python-dateutil = "^2.8.19.10"
 
+[tool.poetry.extras]
+db = ["sqlmodel", "psycopg2-binary"]
+
 [tool.poetry.group.dev.dependencies]
 black = "^22.3.0"
 isort = "^5.8.0"

tests/io/test_writers.py

@@ -19,7 +19,7 @@
 )
 from nannyml.drift.multivariate.data_reconstruction import DataReconstructionDriftCalculator
 from nannyml.drift.univariate import UnivariateDriftCalculator
-from nannyml.io import DatabaseWriter, PickleFileWriter, RawFilesWriter
+from nannyml.io import PickleFileWriter, RawFilesWriter
 from nannyml.performance_calculation import PerformanceCalculator
 from nannyml.performance_estimation.confidence_based import CBPE
 from nannyml.performance_estimation.direct_loss_estimation import DLE
@@ -303,8 +303,11 @@ def test_raw_files_writer_raises_no_exceptions_when_writing_to_csv(result):  # n
 )
 def test_database_writer_raises_no_exceptions_when_writing(result):  # noqa: D103
     try:
+        from nannyml.io.db import DatabaseWriter
         writer = DatabaseWriter(connection_string='sqlite:///', model_name='test')
         writer.write(result)
+    except ImportError:
+        pytest.skip("`db` module is not available.")
     except Exception as exc:
         pytest.fail(f"an unexpected exception occurred: {exc}")
 
@@ -378,6 +381,7 @@ def test_pickle_file_writer_raises_no_exceptions_when_writing(result):  # noqa:
 )
 def test_database_writer_exports_correctly(result, table_name, expected_row_count):  # noqa: D103
     try:
+        from nannyml.io.db import DatabaseWriter
         writer = DatabaseWriter(connection_string='sqlite:///test.db', model_name='test')
         writer.write(result)
 
@@ -386,12 +390,15 @@ def test_database_writer_exports_correctly(result, table_name, expected_row_coun
         with sqlite3.connect("test.db", uri=True) as db:
             res = db.cursor().execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()
             assert res[0] == expected_row_count
+    except ImportError:
+        pytest.skip("`db` module is not available.")
 
     except Exception as exc:
         pytest.fail(f"an unexpected exception occurred: {exc}")
 
     finally:
-        os.remove('test.db')
+        if os.path.exists('test.db'):
+            os.remove('test.db')
 
 
 @pytest.mark.parametrize(
@@ -403,6 +410,7 @@ def test_database_writer_exports_correctly(result, table_name, expected_row_coun
 )
 def test_database_writer_deals_with_metric_components(result, table_name):  # noqa: D103
     try:
+        from nannyml.io.db import DatabaseWriter
         writer = DatabaseWriter(connection_string='sqlite:///test.db', model_name='test')
         writer.write(result.filter(metrics=['confusion_matrix']))
 
@@ -416,9 +424,12 @@ def test_database_writer_deals_with_metric_components(result, table_name):  # no
             assert 'false_positive' in sut
             assert 'true_negative' in sut
             assert 'false_negative' in sut
+    except ImportError:
+        pytest.skip("`db` module is not available.")
 
     except Exception as exc:
         pytest.fail(f"an unexpected exception occurred: {exc}")
 
     finally:
-        os.remove('test.db')
+        if os.path.exists('test.db'):
+            os.remove('test.db')

tests/test_runner.py

@@ -15,7 +15,7 @@
     load_synthetic_car_price_dataset,
     load_synthetic_multiclass_classification_dataset,
 )
-from nannyml.io import DatabaseWriter, RawFilesWriter
+from nannyml.io import RawFilesWriter
 from nannyml.io.store import FilesystemStore
 from nannyml.runner import run
 
@@ -140,6 +140,7 @@ def test_runner_executes_for_binary_classification_with_database_writer_without_
     analysis_with_targets = analysis.merge(analysis_targets, on='identifier')
 
     try:
+        from nannyml.io.db import DatabaseWriter
         with tempfile.TemporaryDirectory() as tmpdir:
             run(
                 reference_data=reference,
@@ -166,6 +167,8 @@ def test_runner_executes_for_binary_classification_with_database_writer_without_
                 run_in_console=False,
                 ignore_errors=False,
             )
+    except ImportError:
+        pytest.skip("the `sqlmodel` module is not available")
     except Exception as exc:
         pytest.fail(f"an unexpected exception occurred: {exc}")
 
@@ -176,6 +179,7 @@ def test_runner_executes_for_multiclass_classification_with_database_writer_with
     analysis_with_targets = analysis.merge(analysis_targets, left_index=True, right_index=True)
 
     try:
+        from nannyml.io.db import DatabaseWriter
         with tempfile.TemporaryDirectory() as tmpdir:
             run(
                 reference_data=reference,
@@ -206,6 +210,8 @@ def test_runner_executes_for_multiclass_classification_with_database_writer_with
                 run_in_console=False,
                 ignore_errors=False,
             )
+    except ImportError:
+        pytest.skip("the `sqlmodel` module is not available")
     except Exception as exc:
         pytest.fail(f"an unexpected exception occurred: {exc}")
 
@@ -216,6 +222,7 @@ def test_runner_executes_for_regression_with_database_writer_without_exceptions(
     analysis_with_targets = analysis.join(analysis_targets)
 
     try:
+        from nannyml.io.db import DatabaseWriter
         with tempfile.TemporaryDirectory() as tmpdir:
             run(
                 reference_data=reference,
@@ -241,5 +248,7 @@ def test_runner_executes_for_regression_with_database_writer_without_exceptions(
                 run_in_console=False,
                 ignore_errors=False,
             )
+    except ImportError:
+        pytest.skip("the `sqlmodel` module is not available")
     except Exception as exc:
         pytest.fail(f"an unexpected exception occurred: {exc}")