Repository: topoteretes/cognee
File: cognee/api/health.py (Python)

Code before:
        """Check relational database health."""
        start_time = time.time()
        try:
            from cognee.infrastructure.databases.relational.get_relational_engine import (
                get_relational_engine,
            )
            from cognee.infrastructure.databases.relational.config import get_relational_config

            config = get_relational_config()
            engine = get_relational_engine()

            # Test connection by creating a session
            session = await engine.get_session()
            if session:
                await session.close()

            response_time = int((time.time() - start_time) * 1000)
            return ComponentHealth(
                status=HealthStatus.HEALTHY,
                provider=config.db_provider,
                response_time_ms=response_time,
                details="Connection successful",
            )
        except Exception as e:
            response_time = int((time.time() - start_time) * 1000)
            return ComponentHealth(
                status=HealthStatus.UNHEALTHY,
                provider="unknown",
                response_time_ms=response_time,
                details=f"Connection failed: {str(e)}",
            )

    async def check_vector_db(self) -> ComponentHealth:
        """Check vector database health."""
        start_time = time.time()
        try:
            from cognee.infrastructure.databases.vector.get_vector_engine import get_vector_engine
            from cognee.infrastructure.databases.vector.config import get_vectordb_config

            config = get_vectordb_config()
            engine = get_vector_engine()

            # Test basic operation - just check if engine is accessible
            if hasattr(engine, "health_check"):
                await engine.health_check()
            elif hasattr(engine, "list_tables"):
                # For LanceDB and similar
                engine.list_tables()

            response_time = int((time.time() - start_time) * 1000)
            return ComponentHealth(

Review comment (bug):
Let's log the stacktrace, as we're currently missing out on it if anything breaks

```suggestion
            response_time = int((time.time() - start_time) * 1000)
            logger.error(f"Relational DB health check failed: {str(e)}", exc_info=True)
            return ComponentHealth(
```

Code after:
        """Check relational database health."""
        start_time = time.time()
        try:
            from cognee.infrastructure.databases.relational.get_relational_engine import (
                get_relational_engine,
            )
            from cognee.infrastructure.databases.relational.config import get_relational_config

            config = get_relational_config()
            engine = get_relational_engine()

            # Test connection by creating a session
            session = engine.get_session()
            if session:
                await session.close()

            response_time = int((time.time() - start_time) * 1000)
            return ComponentHealth(
                status=HealthStatus.HEALTHY,
                provider=config.db_provider,
                response_time_ms=response_time,
                details="Connection successful",
            )
        except Exception as e:
            response_time = int((time.time() - start_time) * 1000)
            logger.error(f"Relational DB health check failed: {str(e)}", exc_info=True)
            return ComponentHealth(
                status=HealthStatus.UNHEALTHY,
                provider="unknown",
                response_time_ms=response_time,
                details=f"Connection failed: {str(e)}",
            )

    async def check_vector_db(self) -> ComponentHealth:
        """Check vector database health."""
        start_time = time.time()
        try:
            from cognee.infrastructure.databases.vector.get_vector_engine import get_vector_engine
            from cognee.infrastructure.databases.vector.config import get_vectordb_config

            config = get_vectordb_config()
            engine = get_vector_engine()

            # Test basic operation - just check if engine is accessible
            if hasattr(engine, "health_check"):
                await engine.health_check()
            elif hasattr(engine, "list_tables"):
                # For LanceDB and similar
                engine.list_tables()

            response_time = int((time.time() - start_time) * 1000)