Fixed the CTE validation issue

docs/installation.md

@@ -179,7 +179,7 @@ Start Claude Desktop and the OMCP server should automatically be available for u
 
 ## Integrating with Localhost models 🦙
 
-### Step 1: Install Ollama 
+### Step 1: Install Ollama
 Download and install [Ollama](https://ollama.com/) from the official website. To check if Ollama has been installed properly, open a terminal and type:
 
 ```bash
@@ -195,7 +195,7 @@ Go to the [Ollama models](https://ollama.com/search) and copy the name of the mo
 ollama pull cogito:14b
 ```
 
-The process will take a while depending on the size of the model, but when it finishes type in the terminal: 
+The process will take a while depending on the size of the model, but when it finishes type in the terminal:
 
 ```bash
 ollama list
@@ -207,7 +207,7 @@ if everything went well, you should see the model you have pulled from Ollama. I
 We are going to use [Librechat](https://www.librechat.ai/) as the end-user interface.
 
 1. In the OMCP project, navigate to the directory where the `main.py` file is located, go to the function `def main()` and change the `transport` from `stdio` to `sse`.
-   
+
     ```python
     def main():
     """Main function to run the MCP server."""
@@ -218,11 +218,11 @@ We are going to use [Librechat](https://www.librechat.ai/) as the end-user inter
     ```
 
 2. In the same directory where `main.py` is located, run the following command:
-    
+
     ```python
     python main.py
     ```
-  
+
     You should see something like this in the terminal:
     ```
     INFO:     Started server process [96250]
@@ -272,11 +272,11 @@ We are going to use [Librechat](https://www.librechat.ai/) as the end-user inter
     [+] Running 5/5
      ✔ Container vectordb          Started
      ✔ Container chat-meilisearch  Started
-     ✔ Container chat-mongodb      Started 
+     ✔ Container chat-mongodb      Started
      ✔ Container rag_api           Started
      ✔ Container LibreChat         Started
     ```
-  
+
 9. Finally, go to the browser and type `localhost:3080`, if it is the first time using Librechat, you need to create an account. Then select the model you pulled, in my case `cogito:14b` and in the chat, just next to the `Code Interpreter` you should see the MCP Tool, click on it and select `omop_mcp`.
 
     You should see something like this:

src/omcp/sql_validator.py

@@ -6,6 +6,7 @@
 import sqlglot.expressions as exp
 import typing as t
 import omcp.exceptions as ex
+from sqlglot.optimizer.scope import build_scope
 
 OMOP_TABLES = [
     "care_site",
@@ -91,12 +92,31 @@ def _check_is_select_query(
                 "Only SELECT statements are allowed for security reasons."
             )
 
-    def _check_is_omop_table(self, tables: t.List[exp.Table]) -> ex.TableNotFoundError:
+    def _check_is_omop_table(self, parsed_sql: exp.Expression) -> ex.TableNotFoundError:
+        """
+        Check if all real table references in the query are OMOP CDM tables and
+        ignores CTEs (defined in WITH clauses).
+
+        Args:
+            parsed_sql (exp.Expression): The parsed SQL expression.
+
+        Return:
+            TableNotFoundError: If any non-OMOP tables are found.
+        """
+        root = build_scope(parsed_sql)
+        tables = [
+            source
+            for scope in root.traverse()
+            for alias, (node, source) in scope.selected_sources.items()
+            if isinstance(source, exp.Table)
+        ]
+
         not_omop_tables = [
             table.name.lower()
             for table in tables
             if table.name.lower() not in OMOP_TABLES
         ]
+
         if not_omop_tables:
             return ex.TableNotFoundError(
                 f"Tables not found in OMOP CDM: {', '.join(not_omop_tables)}"
@@ -216,7 +236,7 @@ def validate_sql(self, sql: str):
                 errors.append(ex.ColumnNotFoundError("No columns found in the query."))
 
             # Check is OMOP table
-            errors.append(self._check_is_omop_table(tables))
+            errors.append(self._check_is_omop_table(parsed_sql))
 
             # Check for excluded tables
             errors.append(self._check_unauthorized_tables(tables))

tests/test_sql_validator.py

@@ -125,3 +125,38 @@ def test_source_concept_id_columns(self, validator):
         assert len(errors) == 1
         assert isinstance(errors[0], ex.UnauthorizedColumnError)
         assert "Source value columns are not allowed" in str(errors[0])
+
+    def test_check_is_omop_table_ignores_cte(self, validator):
+        """Test that _check_is_omop_table ignores CTEs"""
+        sql = """
+        WITH patient AS (
+            SELECT person_id, gender_concept_id FROM person
+        ),
+        visits AS (
+            SELECT visit_occurrence_id, person_id FROM visit_occurrence
+        )
+        SELECT p.person_id, v.visit_occurrence_id
+        FROM patient p
+        JOIN visits v ON p.person_id = v.person_id
+        """
+        errors = validator.validate_sql(sql)
+        assert len(errors) == 0, f"Expected no errors, got: {errors}"
+
+    def test_check_is_omop_table_ignores_multiple_ctes(self, validator):
+        """Test that _check_is_omop_table ignores multiple CTEs with non-OMOP tables"""
+        sql = """
+        WITH temp_users AS (
+            SELECT person_id, year_of_birth, gender_concept_id
+            FROM person
+        ),
+        temp_visits AS (
+            SELECT visit_occurrence_id, person_id, visit_start_date
+            FROM visit_occurrence
+        )
+        SELECT p.person_id, p.year_of_birth, v.visit_start_date
+        FROM temp_users p
+        JOIN temp_visits v ON p.person_id = v.person_id
+        """
+
+        errors = validator.validate_sql(sql)
+        assert len(errors) == 0, f"Expected no errors, got: {errors}"