kamidso-back/tests/kamidso/test_mongodb_connector.py

import json

import pytest

from kamidso.engine.ext.datasource.conn_mongodb import (
    FIND_CAP,
    MongoDBConnector,
    referenced_collection_fields,
)


class _Cursor:
    def __init__(self, docs):
        self._docs = docs

    def limit(self, n):
        return self._docs[:n]


class _Collection:
    def __init__(self, *, docs=None, count=None, agg=None, field_names=None):
        self.docs = docs or []
        self.count_value = count
        self.agg_docs = agg or []
        self.field_names = list(field_names or [])
        self.count_filters = []
        self.pipelines = []

    def find(self, filt=None, projection=None, **_kwargs):
        return _Cursor(self.docs)

    def count_documents(self, filt, **_kwargs):
        self.count_filters.append(filt)
        if self.count_value is not None:
            return self.count_value
        return len(self.docs)

    def aggregate(self, pipeline, **_kwargs):
        self.pipelines.append(pipeline)
        if self.field_names and any(
            isinstance(stage, dict) and "$project" in stage for stage in pipeline
        ):
            return [{"keys": list(self.field_names)}]
        return list(self.agg_docs)


class _Database:
    def __init__(self, collection):
        self._collection = collection

    def __getitem__(self, name):
        return self._collection


class _Client:
    def __init__(self, database):
        self._database = database

    def __getitem__(self, name):
        return self._database

    def close(self):
        return None


def _connector(collection: _Collection) -> MongoDBConnector:
    return MongoDBConnector(_Client(_Database(collection)), "app")


def test_count_runs_on_server_not_find_cap():
    coll = _Collection(count=5_000_000)
    result = _connector(coll).run('{"collection":"users","count":true}')
    assert result == [["count"], [5_000_000]]
    assert coll.count_filters == [{}]


def test_count_keeps_filter_and_zero():
    coll = _Collection(count=0)
    result = _connector(coll).run(
        '{"collection":"users","count":true,"filter":{"role":"admin"}}'
    )
    assert result == [["count"], [0]]
    assert coll.count_filters == [{"role": "admin"}]


def test_find_is_capped_and_is_not_a_count():
    coll = _Collection(docs=[{"_id": i} for i in range(500)])
    result = _connector(coll).run('{"collection":"users","limit":5000}')
    assert result[0] == ["_id"]
    assert len(result) == FIND_CAP + 1


def test_aggregate_group_and_rejects_out():
    coll = _Collection(agg=[{"_id": "admin", "n": 12}])
    result = _connector(coll).run(
        json.dumps(
            {
                "collection": "users",
                "aggregate": [{"$group": {"_id": "$role", "n": {"$sum": 1}}}],
            }
        )
    )
    assert result == [["_id", "n"], ["admin", 12]]
    with pytest.raises(ValueError, match="not allowed"):
        _connector(coll).run(
            '{"collection":"users","aggregate":[{"$out":"dump"}]}'
        )


def test_aggregate_rejects_full_collection_dump():
    coll = _Collection(agg=[{"created_at": "2026-05-04"}])
    with pytest.raises(ValueError, match="must reduce"):
        _connector(coll).run(
            json.dumps(
                {
                    "collection": "users",
                    "aggregate": [
                        {"$project": {"created_at": 1}},
                        {"$sort": {"created_at": 1}},
                    ],
                }
            )
        )
    assert coll.pipelines == []


def test_aggregate_allows_schema_sample_limit():
    coll = _Collection(agg=[{"_id": "x"}])
    result = _connector(coll).run(
        json.dumps(
            {
                "collection": "users",
                "aggregate": [{"$sort": {"created_at": -1}}, {"$limit": 5}],
            }
        )
    )
    assert result[0] == ["_id"]
    assert coll.pipelines


def test_aggregate_allows_sample_limit_25():
    coll = _Collection(agg=[{"_id": "x"}])
    result = _connector(coll).run(
        json.dumps(
            {
                "collection": "users",
                "aggregate": [{"$limit": 25}],
            }
        )
    )
    assert result[0] == ["_id"]
    assert coll.pipelines


def test_aggregate_rejects_large_limit_dump():
    coll = _Collection(agg=[{"_id": "x"}])
    with pytest.raises(ValueError, match="must reduce"):
        _connector(coll).run(
            json.dumps({"collection": "users", "aggregate": [{"$limit": 200}]})
        )
    assert coll.pipelines == []


def test_group_sort_alias_is_not_an_unknown_field():
    coll = _Collection(
        agg=[{"_id": "open", "count": 12}],
        field_names=["_id", "status", "created_at"],
    )
    result = _connector(coll).run(
        json.dumps(
            {
                "collection": "tickets",
                "aggregate": [
                    {"$group": {"_id": "$status", "count": {"$sum": 1}}},
                    {"$sort": {"count": -1}},
                    {"$limit": 25},
                ],
            }
        )
    )
    assert result == [["_id", "count"], ["open", 12]]
    assert coll.pipelines


def test_group_match_alias_is_not_an_unknown_field():
    coll = _Collection(
        agg=[{"_id": "open", "n": 12}],
        field_names=["_id", "status", "created_at"],
    )
    result = _connector(coll).run(
        json.dumps(
            {
                "collection": "tickets",
                "aggregate": [
                    {"$group": {"_id": "$status", "n": {"$sum": 1}}},
                    {"$match": {"n": {"$gte": 1}}},
                ],
            }
        )
    )
    assert result == [["_id", "n"], ["open", 12]]


def test_referenced_fields_skip_group_aliases():
    fields = referenced_collection_fields(
        {
            "aggregate": [
                {"$group": {"_id": "$status", "count": {"$sum": 1}}},
                {"$sort": {"count": -1}},
                {"$match": {"count": {"$gte": 1}}},
            ]
        }
    )
    assert "count" not in fields
    assert "status" in fields


def test_lookup_without_group_is_rejected():
    coll = _Collection(agg=[{"user_id": "1"}])
    with pytest.raises(ValueError, match=r"\$lookup"):
        _connector(coll).run(
            json.dumps(
                {
                    "collection": "users",
                    "aggregate": [
                        {
                            "$lookup": {
                                "from": "wallet_transactions",
                                "localField": "_id",
                                "foreignField": "user_id",
                                "as": "tx",
                            }
                        }
                    ],
                }
            )
        )


def test_lookup_then_group_is_allowed():
    coll = _Collection(agg=[{"_id": "u1", "n": 3}])
    result = _connector(coll).run(
        json.dumps(
            {
                "collection": "users",
                "aggregate": [
                    {
                        "$lookup": {
                            "from": "wallet_transactions",
                            "localField": "_id",
                            "foreignField": "user_id",
                            "as": "tx",
                        }
                    },
                    {"$group": {"_id": "$_id", "n": {"$sum": 1}}},
                ],
            }
        )
    )
    assert result == [["_id", "n"], ["u1", 3]]


def test_unknown_field_lists_real_keys():
    coll = _Collection(count=10, field_names=["_id", "created_at", "role"])
    with pytest.raises(ValueError, match="Unknown field") as excinfo:
        _connector(coll).run(
            '{"collection":"users","count":true,"filter":{"acquisition_source":"ads"}}'
        )
    msg = str(excinfo.value)
    assert "acquisition_source" in msg
    assert "created_at" in msg
    assert "role" in msg
    assert "Do not guess" in msg


def test_unknown_aggregate_field_lists_real_keys():
    coll = _Collection(field_names=["_id", "created_at", "role"])
    with pytest.raises(ValueError, match="Unknown field") as excinfo:
        _connector(coll).run(
            json.dumps(
                {
                    "collection": "users",
                    "aggregate": [
                        {"$group": {"_id": "$referral_code", "n": {"$sum": 1}}}
                    ],
                }
            )
        )
    msg = str(excinfo.value)
    assert "referral_code" in msg
    assert "created_at" in msg
    assert all("$group" not in json.dumps(p) for p in coll.pipelines)


def test_sql_is_rejected():
    with pytest.raises(ValueError, match="does not run SQL"):
        _connector(_Collection()).run("SELECT COUNT(*) FROM users")


def test_mongodb_skips_sql_schema_summary():
    from kamidso.engine.serve.datasource.service.db_summary_client import (
        DBSummaryClient,
        supports_schema_summary,
    )

    assert supports_schema_summary("mysql")
    assert supports_schema_summary("neo4j")
    assert not supports_schema_summary("mongodb")
    assert not supports_schema_summary("Mongo")

    client = object.__new__(DBSummaryClient)
    client.db_summary_embedding("talaban", "mongodb")
    with pytest.raises(ValueError, match="table-schema summary"):
        DBSummaryClient.create_summary_client("talaban", "mongodb")