Fix/run sql api key (#78)

* wip

* release: bump version to 0.4.3 (patch release)

* create changelog file on new release

* fix run_sql when an API is passed

Author: Rodriguespn

changelog/0.4.3.md

@@ -0,0 +1,5 @@
+# [0.4.3] - 2025-07-23
+
+## Fixed
+
+- `run_sql` for virtual workspaces when an API key is given.

scripts/mark-release.sh

@@ -96,15 +96,35 @@ git commit -m "release: bump version to $NEW_VERSION ($RELEASE_TYPE release)"
 git push origin "$CURRENT_BRANCH"
 
 echo ""
+# Create changelog file if it doesn't exist
+CHANGELOG_FILE="changelog/$NEW_VERSION.md"
+if [ ! -f "$CHANGELOG_FILE" ]; then
+    CURRENT_DATE=$(date +%Y-%m-%d)
+    echo "# [$NEW_VERSION] - $CURRENT_DATE
+
+## Added
+-
+-
+
+## Fixed
+-
+- " > "$CHANGELOG_FILE"
+
+    # Add changelog to git
+    git add "$CHANGELOG_FILE"
+    git commit --amend --no-edit
+fi
+
 echo -e "${GREEN}✅ Branch marked for release!${NC}"
 echo ""
 echo -e "${BLUE}Next steps:${NC}"
-echo "1. Push your branch: ${YELLOW}git push origin $CURRENT_BRANCH${NC}"
-echo "2. Create PR to main"
-echo "3. When PR is merged, automatic PyPI publication will be triggered"
-echo "4. Optionally create git tag manually: ${YELLOW}git tag v$NEW_VERSION && git push origin v$NEW_VERSION${NC}"
+echo "1. Update the changelog at: ${YELLOW}$CHANGELOG_FILE${NC}"
+echo "2. Push your branch: ${YELLOW}git push origin $CURRENT_BRANCH${NC}"
+echo "3. Create PR to main"
+echo "4. When PR is merged, automatic PyPI publication will be triggered"
+echo "5. Optionally create git tag manually: ${YELLOW}git tag v$NEW_VERSION && git push origin v$NEW_VERSION${NC}"
 echo ""
 echo -e "${BLUE}Release details:${NC}"
 echo "  🏷️  Release type: $RELEASE_TYPE"
 echo "  📦 New version: $NEW_VERSION"
-echo "  📝 Version file updated and committed"
+echo "  📝 Version file and changelog updated and committed"

src/api/tools/tools.py

@@ -23,6 +23,10 @@
     query_graphql_organizations,
 )
 from src.api.tools.types import Tool, WorkspaceTarget
+from src.auth.session_credentials_manager import (
+    get_session_credentials_manager,
+    invalidate_credentials,
+)
 from src.utils.uuid_validation import validate_workspace_id, validate_uuid_string
 from src.utils.elicitation import try_elicitation, ElicitationError
 from src.logger import get_logger
@@ -35,6 +39,102 @@
 )
 
 
+class DatabaseCredentials(BaseModel):
+    """Schema for database authentication credentials when using API key."""
+
+    username: str = Field(..., description="Database username for authentication")
+    password: str = Field(..., description="Database password for authentication")
+
+
+async def _get_database_credentials(
+    ctx: Context, target: WorkspaceTarget, database_name: str | None = None
+) -> tuple[str, str]:
+    """
+    Get database credentials based on the authentication method.
+
+    Args:
+        ctx: The MCP context
+        target: The workspace target
+        database_name: The database name to use for key generation
+
+    Returns:
+        Tuple of (username, password)
+
+    Raises:
+        Exception: If credentials cannot be obtained
+    """
+    settings = config.get_settings()
+
+    # Check if we're using API key authentication
+    is_using_api_key = (
+        not settings.is_remote
+        and isinstance(settings, config.LocalSettings)
+        and settings.api_key is not None
+    )
+
+    if is_using_api_key:
+        # For API key authentication, we need database credentials
+        # Generate database key using credentials manager
+        credentials_manager = get_session_credentials_manager()
+        database_key = credentials_manager.generate_database_key(
+            workspace_name=target.name, database_name=database_name
+        )
+
+        # Check if we have cached credentials for this database
+        if credentials_manager.has_credentials(database_key):
+            cached_creds = credentials_manager.get_credentials(database_key)
+            if cached_creds:
+                logger.debug(f"Using cached credentials for workspace: {target.name}")
+                return cached_creds
+
+        # Dedicated workspaces: need to request database credentials from user
+        elicitation_message = (
+            f"API key authentication detected. To connect to the dedicated workspace '{target.name}', "
+            f"please provide your database username and password for this workspace."
+        )
+
+        try:
+            elicitation_result, error = await try_elicitation(
+                ctx=ctx, message=elicitation_message, schema=DatabaseCredentials
+            )
+
+            if error == ElicitationError.NOT_SUPPORTED:
+                # Fallback: raise exception with clear message
+                raise Exception(
+                    "Database credentials required for API key authentication on dedicated workspaces. "
+                    f"Please provide your database username and password for workspace '{target.name}'. "
+                    "You can obtain these credentials from your SingleStore portal. "
+                    "Note: This is different from your SingleStore account credentials - these are "
+                    "database-specific credentials for connecting to the workspace."
+                )
+            elif elicitation_result.status == "success" and elicitation_result.data:
+                username = elicitation_result.data.username
+                password = elicitation_result.data.password
+
+                # Store credentials in session cache for future use
+                try:
+                    credentials_manager.store_credentials(
+                        database_key, username, password
+                    )
+                    logger.debug(f"Cached credentials for workspace: {target.name}")
+                except Exception as e:
+                    logger.warning(f"Failed to cache credentials: {e}")
+
+                return (username, password)
+            else:
+                raise Exception(
+                    "Database credentials are required but were not provided. Please ask the user to provide the database credentials"
+                )
+        except Exception as e:
+            if "Database credentials required" in str(e):
+                raise  # Re-raise our specific credential error
+            logger.error(f"Error during credential elicitation: {e}")
+            raise Exception(f"Failed to obtain database credentials: {str(e)}")
+    else:
+        # JWT authentication: use user_id and access_token as before
+        return __get_user_id(), get_access_token()
+
+
 async def __execute_sql_unified(
     ctx: Context,
     target: WorkspaceTarget,
@@ -60,39 +160,71 @@ async def __execute_sql_unified(
         host = endpoint
         port = None
 
-    s2_manager = S2Manager(
-        host=host,
-        port=port,
-        user=username,
-        password=password,
-        database=database_name,
+    # Generate database key for credential management
+    credentials_manager = get_session_credentials_manager()
+    database_key = credentials_manager.generate_database_key(
+        workspace_name=target.name, database_name=database_name
     )
 
-    workspace_type = "shared/virtual" if target.is_shared else "dedicated"
-    await ctx.info(
-        f"Executing SQL query on {workspace_type} workspace '{target.name}' with database '{database_name}': {sql_query}"
-        "This query may take some time depending on the complexity and size of the data."
-    )
-    s2_manager.execute(sql_query)
-    columns = (
-        [desc[0] for desc in s2_manager.cursor.description]
-        if s2_manager.cursor.description
-        else []
-    )
-    rows = s2_manager.fetchmany()
-    results = []
-    for row in rows:
-        result_dict = {}
-        for i, column in enumerate(columns):
-            result_dict[column] = row[i]
-        results.append(result_dict)
-    s2_manager.close()
-    return {
-        "data": results,
-        "row_count": len(rows),
-        "columns": columns,
-        "status": "Success",
-    }
+    try:
+        s2_manager = S2Manager(
+            host=host,
+            port=port,
+            user=username,
+            password=password,
+            database=database_name,
+        )
+
+        workspace_type = "shared/virtual" if target.is_shared else "dedicated"
+        await ctx.info(
+            f"Executing SQL query on {workspace_type} workspace '{target.name}' with database '{database_name}': {sql_query}"
+            "This query may take some time depending on the complexity and size of the data."
+        )
+        s2_manager.execute(sql_query)
+        columns = (
+            [desc[0] for desc in s2_manager.cursor.description]
+            if s2_manager.cursor.description
+            else []
+        )
+        rows = s2_manager.fetchmany()
+        results = []
+        for row in rows:
+            result_dict = {}
+            for i, column in enumerate(columns):
+                result_dict[column] = row[i]
+            results.append(result_dict)
+        s2_manager.close()
+        return {
+            "data": results,
+            "row_count": len(rows),
+            "columns": columns,
+            "status": "Success",
+        }
+    except Exception as e:
+        # Check if this is an authentication error
+        error_msg = str(e).lower()
+        is_auth_error = any(
+            auth_keyword in error_msg
+            for auth_keyword in [
+                "access denied",
+                "authentication failed",
+                "invalid credentials",
+                "login failed",
+                "permission denied",
+                "unauthorized",
+                "auth",
+            ]
+        )
+
+        if is_auth_error:
+            logger.warning(
+                f"Authentication failed for database {database_key}, invalidating cached credentials"
+            )
+            invalidate_credentials(database_key)
+            raise Exception(f"Authentication failed: {str(e)}")
+        else:
+            # Non-authentication error, re-raise as-is
+            raise
 
 
 def __get_virtual_workspace(virtual_workspace_id: str):
@@ -781,19 +913,62 @@ async def run_sql(
     if target.is_shared and target.database_name and not database_name:
         database_name = target.database_name
 
-    username = __get_user_id()
-    password = get_access_token()
+    # Get database credentials based on authentication method
+    try:
+        username, password = await _get_database_credentials(ctx, target, database_name)
+    except Exception as e:
+        if "Database credentials required" in str(e):
+            # Handle the specific case where elicitation is not supported
+            return {
+                "status": "error",
+                "message": str(e),
+                "error_code": "CREDENTIALS_REQUIRED",
+                "workspace_id": validated_id,
+                "workspace_name": target.name,
+                "workspace_type": "shared" if target.is_shared else "dedicated",
+                "instruction": (
+                    "Please call this function again with the same parameters once you have "
+                    "the database credentials available, or ask the user to provide their "
+                    "database username and password for this workspace."
+                ),
+            }
+        else:
+            return {
+                "status": "error",
+                "message": f"Failed to obtain database credentials: {str(e)}",
+                "error_code": "AUTHENTICATION_ERROR",
+            }
 
     # Execute the SQL query
     start_time = time.time()
-    result = await __execute_sql_unified(
-        ctx=ctx,
-        target=target,
-        sql_query=sql_query,
-        username=username,
-        password=password,
-        database=database_name,
-    )
+    try:
+        result = await __execute_sql_unified(
+            ctx=ctx,
+            target=target,
+            sql_query=sql_query,
+            username=username,
+            password=password,
+            database=database_name,
+        )
+    except Exception as e:
+        # Check if this is an authentication error from __execute_sql_unified
+        if "Authentication failed:" in str(e):
+            # Authentication error already handled by __execute_sql_unified (credentials invalidated)
+            return {
+                "status": "error",
+                "message": str(e),
+                "error_code": "AUTHENTICATION_ERROR",
+                "workspace_id": validated_id,
+                "workspace_name": target.name,
+                "workspace_type": "shared" if target.is_shared else "dedicated",
+                "instruction": (
+                    "Authentication failed. Please provide valid database credentials "
+                    "for this workspace and try again."
+                ),
+            }
+        else:
+            # Non-authentication error, re-raise
+            raise
 
     results_data = result.get("data", [])