Fix BUG-005: Trace Context & Misc Regressions

nadeem4 · nadeem4 · commit d05ea6626e6d · 2026-01-12T16:02:37.000-06:00
- Implemented BUG-005: Distributed Tracing (TraceContextFilter, TraceID in GraphState, traced_node wrapper)
- Fixed crash in 'nl2sql setup --demo' (DatasourceRegistry misuse)
- Fixed syntax error in Intent Validator prompt (unescaped JSON braces)
- Fixed SqliteAdapter DryRunResult schema mismatch (valid -&gt; is_valid)
- Updated all Adapters (Postgres, MySQL, MSSQL) get_dialect to return string names
- Added unit tests for tracing
diff --git a/audit/remediation_plan.md b/audit/remediation_plan.md
@@ -28,10 +28,11 @@ This document serves as the master backlog for addressing findings from the Arch
   - **Fix**: Implemented idempotent `refresh_schema` and `refresh_examples` in `OrchestratorVectorStore`, along with dynamic `register_datasource` in `DatasourceRegistry`.
   - **Status**: Fixed. Unit tests added in `tests/unit/test_schema_lifecycle.py`.
 
-- [ ] **BUG-005: Missing Distributed Tracing** (High)
+- [x] **BUG-005: Missing Distributed Tracing** (High)
   - **Component**: Observability / Logging
   - **Issue**: Logs lack a unique `trace_id` per request, making concurrent request debugging impossible in multi-threaded environments.
-  - **Fix**: Inject `trace_id` at the `GraphState` entry point and propagate it to `Python LogRecord` context.
+  - **Fix**: Implemented `trace_id` in `GraphState`, `TraceContextFilter` in logger, and `traced_node` wrapper for context propagation.
+  - **Status**: Fixed. Verified in `tests/unit/test_tracing.py`.
 
 ## 🟡 Medium & Low Priority Bugs
 
diff --git a/configs/llm.demo.yaml b/configs/llm.demo.yaml
@@ -6,9 +6,3 @@ default:
   model: gpt-4o
   temperature: 0.0
   api_key: ${env:OPENAI_API_KEY}
-agents:
-  intent_validator:
-    provider: openai
-    model: gpt-4o-mini
-    temperature: 0.0
-    api_key: ${env:OPENAI_API_KEY}
diff --git a/packages/adapters/mssql/src/nl2sql_mssql/adapter.py b/packages/adapters/mssql/src/nl2sql_mssql/adapter.py
@@ -90,7 +90,7 @@ def explain(self, sql: str) -> QueryPlan:
 
     def get_dialect(self) -> str:
         """MSSQL uses T-SQL dialect."""
-        return mssql.dialect()
+        return mssql.dialect.name
 
     def cost_estimate(self, sql: str) -> CostEstimate:
         import re
diff --git a/packages/adapters/mysql/src/nl2sql_mysql/adapter.py b/packages/adapters/mysql/src/nl2sql_mysql/adapter.py
@@ -116,5 +116,5 @@ def cost_estimate(self, sql: str) -> CostEstimate:
         return CostEstimate(estimated_cost=0.0, estimated_rows=0)
 
     def get_dialect(self) -> str:
-        return mysql.dialect()
+        return mysql.dialect.name
 
diff --git a/packages/adapters/postgres/src/nl2sql_postgres/adapter.py b/packages/adapters/postgres/src/nl2sql_postgres/adapter.py
@@ -108,6 +108,4 @@ def cost_estimate(self, sql: str) -> CostEstimate:
             return CostEstimate(estimated_cost=0.0, estimated_rows=0)
 
     def get_dialect(self) -> str:
-        return postgresql.dialect()
-
-
+        return postgresql.dialect.name
diff --git a/packages/adapters/sqlite/src/nl2sql_sqlite/adapter.py b/packages/adapters/sqlite/src/nl2sql_sqlite/adapter.py
@@ -64,9 +64,9 @@ def dry_run(self, query: str) -> DryRunResult:
         try:
             with self.engine.connect() as conn:
                 conn.execute(text(f"EXPLAIN QUERY PLAN {query}"))
-            return DryRunResult(valid=True, error=None)
+            return DryRunResult(is_valid=True, error_message=None)
         except Exception as e:
-            return DryRunResult(valid=False, error=str(e))
+            return DryRunResult(is_valid=False, error_message=str(e))
 
     def explain(self, query: str) -> QueryPlan:
          return QueryPlan(original_query=query, plan="EXPLAIN QUERY PLAN not fully parsed")
@@ -83,4 +83,4 @@ def cost_estimate(self, query: str) -> CostEstimate:
 
 
     def get_dialect(self) -> str:
-        return sqlite.dialect()
+        return sqlite.dialect.name
diff --git a/packages/cli/src/nl2sql_cli/demo/manager.py b/packages/cli/src/nl2sql_cli/demo/manager.py
@@ -210,6 +210,7 @@ def index_demo_data(self):
         from nl2sql.services.vector_store import OrchestratorVectorStore
         from nl2sql.services.llm import LLMRegistry
         from nl2sql.configs import ConfigManager 
+        from nl2sql.datasources import DatasourceRegistry 
         
         try:
             indexer_config_manager = ConfigManager(self.project_root)
@@ -224,7 +225,8 @@ def index_demo_data(self):
             except Exception:
                 pass # Optional for schema indexing
                 
-            run_indexing(configs, settings.vector_store_path, v_store, llm_registry)
+            registry = DatasourceRegistry(configs)
+            run_indexing(registry, settings.vector_store_path, v_store, llm_registry)
             return True
         except Exception as e:
             self.print_error(f"Indexing Failed: {e}")
diff --git a/packages/core/src/nl2sql/common/logger.py b/packages/core/src/nl2sql/common/logger.py
@@ -1,8 +1,26 @@
 import logging
 import json
 import time
-from typing import Any, Dict
+import contextvars
+from contextlib import contextmanager
+from typing import Any, Dict, Optional
 
+_trace_id_ctx = contextvars.ContextVar("trace_id", default=None)
+
+class TraceContextFilter(logging.Filter):
+    """Injects trace_id from contextvar into the log record."""
+    def filter(self, record):
+        record.trace_id = _trace_id_ctx.get()
+        return True
+
+@contextmanager
+def trace_context(trace_id: str):
+    """Context manager to set the trace_id for the current context."""
+    token = _trace_id_ctx.set(trace_id)
+    try:
+        yield
+    finally:
+        _trace_id_ctx.reset(token)
 
 class JsonFormatter(logging.Formatter):
     """Formatter that outputs JSON strings after parsing the LogRecord."""
@@ -23,13 +41,16 @@ def format(self, record: logging.LogRecord) -> str:
             "message": record.getMessage(),
         }
         
+        if getattr(record, "trace_id", None):
+            log_record["trace_id"] = record.trace_id
+
         # Standard LogRecord attributes to ignore
         standard_attrs = {
             "args", "asctime", "created", "exc_info", "exc_text", "filename",
             "funcName", "levelname", "levelno", "lineno", "module",
             "msecs", "message", "msg", "name", "pathname", "process",
             "processName", "relativeCreated", "stack_info", "thread", "threadName",
-            "taskName"
+            "taskName", "trace_id"
         }
 
         for key, value in record.__dict__.items():
@@ -54,11 +75,19 @@ def configure_logging(level: str = "INFO", json_format: bool = False):
         root_logger.removeHandler(handler)
         
     handler = logging.StreamHandler()
+    handler.addFilter(TraceContextFilter())
     
     if json_format:
         handler.setFormatter(JsonFormatter())
     else:
-        handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
+        # Include trace_id in standard format if present
+        # This is a bit tricky with dynamic formatting, usually easier to check record in formatter
+        # For simplicity, we stick to standard format but maybe prepend trace_id if possible?
+        # We'll stick to a standard format for text logs for now, trace_id mainly for JSON/Production
+        formatter = logging.Formatter(
+            "%(asctime)s - [%(trace_id)s] - %(name)s - %(levelname)s - %(message)s"
+        )
+        handler.setFormatter(formatter)
         
     root_logger.addHandler(handler)
     
diff --git a/packages/core/src/nl2sql/pipeline/graph.py b/packages/core/src/nl2sql/pipeline/graph.py
@@ -15,11 +15,29 @@
 from nl2sql.services.vector_store import OrchestratorVectorStore
 from nl2sql.services.llm import LLMRegistry
 from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.common.logger import trace_context
 
 
 LLMCallable = Union[Callable[[str], str], Runnable]
 
 
+def traced_node(node: Callable):
+    """Wraps a node to inject trace_id from state into the logging context."""
+    def wrapper(state: Union[Dict, Any]):
+        # Extract trace_id from state (dict or object)
+        tid = None
+        if isinstance(state, dict):
+            tid = state.get("trace_id")
+        else:
+            tid = getattr(state, "trace_id", None)
+            
+        if tid:
+            with trace_context(tid):
+                return node(state)
+        return node(state)
+    return wrapper
+
+
 def build_graph(
     registry: DatasourceRegistry,
     llm_registry: LLMRegistry,
@@ -112,12 +130,12 @@ def report_missing_datasource(state: Dict):
             "intermediate_results": [message],
         }
 
-    graph.add_node("semantic_analysis", semantic_node)
-    graph.add_node("intent_validator", intent_validator_node)
-    graph.add_node("decomposer", decomposer_node)
-    graph.add_node("execution_branch", execution_wrapper)
-    graph.add_node("report_missing_datasource", report_missing_datasource)
-    graph.add_node("aggregator", aggregator_node)
+    graph.add_node("semantic_analysis", traced_node(semantic_node))
+    graph.add_node("intent_validator", traced_node(intent_validator_node))
+    graph.add_node("decomposer", traced_node(decomposer_node))
+    graph.add_node("execution_branch", traced_node(execution_wrapper))
+    graph.add_node("report_missing_datasource", traced_node(report_missing_datasource))
+    graph.add_node("aggregator", traced_node(aggregator_node))
 
     graph.set_entry_point("semantic_analysis")
     
@@ -140,6 +158,7 @@ def continue_to_subqueries(state: GraphState):
                 branches.append(Send("report_missing_datasource", {"user_query": sq.query}))
             
             payload = {
+                "trace_id": state.trace_id,
                 "user_query": sq.query,
                 "selected_datasource_id": sq.datasource_id,
                 "complexity": sq.complexity,
diff --git a/packages/core/src/nl2sql/pipeline/nodes/intent_validator/prompts.py b/packages/core/src/nl2sql/pipeline/nodes/intent_validator/prompts.py
@@ -17,11 +17,11 @@
 
 [OUTPUT SCHEMA]
 Return a JSON object matching `IntentValidationResult`:
-{
+{{
   "is_safe": boolean,
   "violation_category": "jailbreak" | "pii_exfiltration" | "destructive" | "system_probing" | "none",
   "reasoning": "string"
-}
+}}
 
 [USER_QUERY]
 {user_query}
diff --git a/packages/core/src/nl2sql/pipeline/state.py b/packages/core/src/nl2sql/pipeline/state.py
@@ -2,6 +2,7 @@
 
 from typing import Any, Dict, List, Optional, Annotated, Literal
 import operator
+import uuid
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -40,6 +41,7 @@ class GraphState(BaseModel):
     """
     model_config = ConfigDict(extra="ignore", arbitrary_types_allowed=True)
 
+    trace_id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Distributed unique trace ID.")
     user_query: str = Field(description="Canonical user query.")
     complexity: Literal["simple", "complex"] = Field(
         default="complex",
diff --git a/packages/core/tests/unit/test_tracing.py b/packages/core/tests/unit/test_tracing.py
@@ -0,0 +1,81 @@
+import pytest
+import logging
+import json
+import uuid
+from unittest.mock import MagicMock
+from nl2sql.common.logger import TraceContextFilter, trace_context, JsonFormatter, _trace_id_ctx
+from nl2sql.pipeline.state import GraphState
+from nl2sql.pipeline.graph import traced_node
+
+def test_trace_context_propagation(capsys):
+    """Verifies that trace_id is correctly injected into logs."""
+    logger = logging.getLogger("test_logger")
+    logger.setLevel(logging.INFO)
+    
+    # Add filter and formatter
+    handler = logging.StreamHandler()
+    handler.addFilter(TraceContextFilter())
+    handler.setFormatter(JsonFormatter())
+    logger.addHandler(handler)
+    
+    trace_id = "test-trace-123"
+    
+    with trace_context(trace_id):
+        logger.info("Test message")
+        
+    captured = capsys.readouterr()
+    log_line = captured.err # logging defaults to stderr usually
+    
+    # Depending on capsys and logging config, it might be in different streams
+    # If no output, we might need to verify record directly
+    # But let's check basic record passing first
+    
+def test_trace_context_filter_logic():
+    """Verifies filter logic directly."""
+    f = TraceContextFilter()
+    record = logging.LogRecord("name", logging.INFO, "path", 1, "msg", {}, None)
+    
+    trace_id = "test-uuid"
+    with trace_context(trace_id):
+        f.filter(record)
+        assert getattr(record, "trace_id") == trace_id
+
+def test_graph_state_trace_id_generation():
+    """Verifies GraphState auto-generates trace_id."""
+    state = GraphState(user_query="test")
+    assert state.trace_id is not None
+    assert isinstance(state.trace_id, str)
+    assert len(state.trace_id) > 10 # UUID length
+
+def test_traced_node_wrapper():
+    """Verifies traced_node wrapper injects context."""
+    
+    mock_func = MagicMock()
+    
+    def actual_func(state):
+        # Inside function, check if context is set
+        tid = _trace_id_ctx.get()
+        mock_func(tid)
+        return {"result": "ok"}
+    
+    wrapped = traced_node(actual_func)
+    
+    test_tid = "trace-456"
+    state = {"trace_id": test_tid, "data": "foo"}
+    
+    wrapped(state)
+    
+    mock_func.assert_called_with(test_tid)
+
+def test_traced_node_wrapper_missing_id():
+    """Verifies wrapper handles missing trace_id safely."""
+    mock_func = MagicMock()
+    def actual_func(state):
+        tid = _trace_id_ctx.get()
+        mock_func(tid)
+        return {}
+        
+    wrapped = traced_node(actual_func)
+    wrapped({})
+    
+    mock_func.assert_called_with(None)
