├── requirements.txt ├── pytest.ini ├── CHANGELOG.md ├── tests ├── README.md └── test_azure_postgresql_mcp.py ├── .github ├── CODE_OF_CONDUCT.md ├── ISSUE_TEMPLATE.md └── PULL_REQUEST_TEMPLATE.md ├── LICENSE.md ├── CONTRIBUTING.md ├── .gitignore ├── README.md └── src └── azure_postgresql_mcp.py /requirements.txt: -------------------------------------------------------------------------------- 1 | azure-identity==1.21.0 2 | azure-mgmt-postgresqlflexibleservers==1.1.0 3 | mcp[cli]==1.6.0 4 | psycopg[binary]==3.2.6 -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | log_file = tests/test_errors.log 3 | log_file_level = ERROR 4 | log_file_format = %(asctime)s - %(levelname)s - %(message)s 5 | log_cli = false 6 | log_cli_level = INFO 7 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | ## [project-title] Changelog 2 | 3 | 4 | # x.y.z (yyyy-mm-dd) 5 | 6 | *Features* 7 | * ... 8 | 9 | *Bug Fixes* 10 | * ... 11 | 12 | *Breaking Changes* 13 | * ... 14 | -------------------------------------------------------------------------------- /tests/README.md: -------------------------------------------------------------------------------- 1 | # Running Tests for Azure PostgreSQL MCP 2 | 3 | To run the tests, first ensure you have `pytest` installed. You can install it using: 4 | 5 | ```bash 6 | pip install pytest 7 | ``` 8 | 9 | To execute the tests, use the following command: 10 | 11 | ```bash 12 | PYTHONPATH=src pytest --color=yes -v 13 | ``` 14 | 15 | Error logs will be written to `tests/test_errors.log` as configured in `pytest.ini`. 16 | -------------------------------------------------------------------------------- /.github/CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Microsoft Open Source Code of Conduct 2 | 3 | This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). 4 | 5 | Resources: 6 | 7 | - [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/) 8 | - [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) 9 | - Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns 10 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | 4 | > Please provide us with the following information: 5 | > --------------------------------------------------------------- 6 | 7 | ### This issue is for a: (mark with an `x`) 8 | ``` 9 | - [ ] bug report -> please search issues before submitting 10 | - [ ] feature request 11 | - [ ] documentation issue or request 12 | - [ ] regression (a behavior that used to work and stopped in a new release) 13 | ``` 14 | 15 | ### Minimal steps to reproduce 16 | > 17 | 18 | ### Any log messages given by the failure 19 | > 20 | 21 | ### Expected/desired behavior 22 | > 23 | 24 | ### OS and Version? 25 | > Windows 7, 8 or 10. Linux (which distribution). macOS (Yosemite? El Capitan? Sierra?) 26 | 27 | ### Versions 28 | > 29 | 30 | ### Mention any other details that might be useful 31 | 32 | > --------------------------------------------------------------- 33 | > Thanks! We'll be in touch soon. 34 | -------------------------------------------------------------------------------- /.github/PULL_REQUEST_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | ## Purpose 2 | 3 | * ... 4 | 5 | ## Does this introduce a breaking change? 6 | 7 | ``` 8 | [ ] Yes 9 | [ ] No 10 | ``` 11 | 12 | ## Pull Request Type 13 | What kind of change does this Pull Request introduce? 14 | 15 | 16 | ``` 17 | [ ] Bugfix 18 | [ ] Feature 19 | [ ] Code style update (formatting, local variables) 20 | [ ] Refactoring (no functional changes, no api changes) 21 | [ ] Documentation content changes 22 | [ ] Other... Please describe: 23 | ``` 24 | 25 | ## How to Test 26 | * Get the code 27 | 28 | ``` 29 | git clone [repo-address] 30 | cd [repo-name] 31 | git checkout [branch-name] 32 | npm install 33 | ``` 34 | 35 | * Test the code 36 | 37 | ``` 38 | ``` 39 | 40 | ## What to Check 41 | Verify that the following are valid 42 | * ... 43 | 44 | ## Other Information 45 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) Microsoft Corporation. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to [project-title] 2 | 3 | This project welcomes contributions and suggestions. Most contributions require you to agree to a 4 | Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us 5 | the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com. 6 | 7 | When you submit a pull request, a CLA bot will automatically determine whether you need to provide 8 | a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions 9 | provided by the bot. You will only need to do this once across all repos using our CLA. 10 | 11 | This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). 12 | For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or 13 | contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments. 14 | 15 | - [Code of Conduct](#coc) 16 | - [Issues and Bugs](#issue) 17 | - [Feature Requests](#feature) 18 | - [Submission Guidelines](#submit) 19 | 20 | ## Code of Conduct 21 | Help us keep this project open and inclusive. Please read and follow our [Code of Conduct](https://opensource.microsoft.com/codeofconduct/). 22 | 23 | ## Found an Issue? 24 | If you find a bug in the source code or a mistake in the documentation, you can help us by 25 | [submitting an issue](#submit-issue) to the GitHub Repository. Even better, you can 26 | [submit a Pull Request](#submit-pr) with a fix. 27 | 28 | ## Want a Feature? 29 | You can *request* a new feature by [submitting an issue](#submit-issue) to the GitHub 30 | Repository. If you would like to *implement* a new feature, please submit an issue with 31 | a proposal for your work first, to be sure that we can use it. 32 | 33 | * **Small Features** can be crafted and directly [submitted as a Pull Request](#submit-pr). 34 | 35 | ## Submission Guidelines 36 | 37 | ### Submitting an Issue 38 | Before you submit an issue, search the archive, maybe your question was already answered. 39 | 40 | If your issue appears to be a bug, and hasn't been reported, open a new issue. 41 | Help us to maximize the effort we can spend fixing issues and adding new 42 | features, by not reporting duplicate issues. Providing the following information will increase the 43 | chances of your issue being dealt with quickly: 44 | 45 | * **Overview of the Issue** - if an error is being thrown a non-minified stack trace helps 46 | * **Version** - what version is affected (e.g. 0.1.2) 47 | * **Motivation for or Use Case** - explain what are you trying to do and why the current behavior is a bug for you 48 | * **Browsers and Operating System** - is this a problem with all browsers? 49 | * **Reproduce the Error** - provide a live example or a unambiguous set of steps 50 | * **Related Issues** - has a similar issue been reported before? 51 | * **Suggest a Fix** - if you can't fix the bug yourself, perhaps you can point to what might be 52 | causing the problem (line of code or commit) 53 | 54 | You can file new issues by providing the above information at the corresponding repository's issues link: https://github.com/[organization-name]/[repository-name]/issues/new]. 55 | 56 | ### Submitting a Pull Request (PR) 57 | Before you submit your Pull Request (PR) consider the following guidelines: 58 | 59 | * Search the repository (https://github.com/[organization-name]/[repository-name]/pulls) for an open or closed PR 60 | that relates to your submission. You don't want to duplicate effort. 61 | 62 | * Make your changes in a new git fork: 63 | 64 | * Commit your changes using a descriptive commit message 65 | * Push your fork to GitHub: 66 | * In GitHub, create a pull request 67 | * If we suggest changes then: 68 | * Make the required updates. 69 | * Rebase your fork and force push to your GitHub repository (this will update your Pull Request): 70 | 71 | ```shell 72 | git rebase master -i 73 | git push -f 74 | ``` 75 | 76 | That's it! Thank you for your contribution! 77 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # UV 98 | # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | #uv.lock 102 | 103 | # poetry 104 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 105 | # This is especially recommended for binary packages to ensure reproducibility, and is more 106 | # commonly ignored for libraries. 107 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 108 | #poetry.lock 109 | 110 | # pdm 111 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 112 | #pdm.lock 113 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 114 | # in version control. 115 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control 116 | .pdm.toml 117 | .pdm-python 118 | .pdm-build/ 119 | 120 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 121 | __pypackages__/ 122 | 123 | # Celery stuff 124 | celerybeat-schedule 125 | celerybeat.pid 126 | 127 | # SageMath parsed files 128 | *.sage.py 129 | 130 | # Environments 131 | .env 132 | .venv 133 | env/ 134 | venv/ 135 | ENV/ 136 | env.bak/ 137 | venv.bak/ 138 | 139 | # Spyder project settings 140 | .spyderproject 141 | .spyproject 142 | 143 | # Rope project settings 144 | .ropeproject 145 | 146 | # mkdocs documentation 147 | /site 148 | 149 | # mypy 150 | .mypy_cache/ 151 | .dmypy.json 152 | dmypy.json 153 | 154 | # Pyre type checker 155 | .pyre/ 156 | 157 | # pytype static type analyzer 158 | .pytype/ 159 | 160 | # Cython debug symbols 161 | cython_debug/ 162 | 163 | # PyCharm 164 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 165 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 166 | # and can be added to the global gitignore or merged into this file. For a more nuclear 167 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 168 | #.idea/ 169 | 170 | # Ruff stuff: 171 | .ruff_cache/ 172 | 173 | # PyPI configuration file 174 | .pypirc 175 | 176 | # Ignore test output logs 177 | tests/test_errors.log 178 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Azure Database for PostgreSQL MCP Server (Preview) 2 | 3 | A [Model Context Protocol (MCP)](https://modelcontextprotocol.io/introduction) Server that let’s your AI models talk to data hosted in Azure Database for PostgreSQL according to the MCP standard! 4 | 5 | By utilizing this server, you can effortlessly connect any AI application that supports MCP to your PostgreSQL flexible server (using either PostgreSQL password-based authentication or Microsoft Entra authentication methods), enabling you to provide your business data as meaningful context in a standardized and secure manner. 6 | 7 | This server exposes the following tools, which can be invoked by MCP Clients in your AI agents, AI applications or tools like Claude Desktop and Visual Studio Code: 8 | 9 | - **List all databases** in your Azure Database for PostgreSQL flexible server instance. 10 | - **List all tables** in a database along with their schema information. 11 | - **Execute read queries** to retrieve data from your database. 12 | - **Insert or update records** in your database. 13 | - **Create a new table or drop an existing table** in your database. 14 | - **List Azure Database for PostgreSQL flexible server configuration**, including its PostgreSQL version, and compute and storage configurations. * 15 | - Retrieve specific **server parameter values.** * 16 | 17 | _*Available when using Microsoft Entra authentication method_ 18 | 19 | ## Getting Started 20 | 21 | ### Prerequisites 22 | 23 | - [Python](https://www.python.org/downloads/) 3.10 or above 24 | - An Azure Database for PostgreSQL flexible server instance with a database containing your business data. For instructions on creating a flexible instance, setting up a database, and connecting to it, please refer to this [quickstart guide](https://learn.microsoft.com/azure/postgresql/flexible-server/quickstart-create-server). 25 | - An MCP Client application or tool such as [Claude Desktop](https://claude.ai/download) or [Visual Studio Code](https://code.visualstudio.com/download). 26 | 27 | ### Installation 28 | 29 | 1. Clone the `azure-postgresql-mcp` repository: 30 | 31 | ``` 32 | git clone https://github.com/Azure-Samples/azure-postgresql-mcp.git 33 | cd azure-postgresql-mcp 34 | ``` 35 | 36 | Alternatively, you can download only the `azure_postgresql_mcp.py` file to your working folder. 37 | 38 | 2. Create a virtual environment: 39 | 40 | Windows cmd.exe: 41 | ``` 42 | python -m venv azure-postgresql-mcp-venv 43 | .\azure-postgresql-mcp-venv\Scripts\activate.bat 44 | ``` 45 | Windows Powershell: 46 | ``` 47 | python -m venv azure-postgresql-mcp-venv 48 | .\azure-postgresql-mcp-venv\Scripts\Activate.ps1 49 | ``` 50 | Linux and MacOS: 51 | ``` 52 | python -m venv azure-postgresql-mcp-venv 53 | source ./azure-postgresql-mcp-venv/bin/activate 54 | ``` 55 | 56 | 4. Install the dependencies: 57 | 58 | ``` 59 | pip install mcp[cli] 60 | pip install psycopg[binary] 61 | pip install azure-mgmt-postgresqlflexibleservers 62 | pip install azure-identity 63 | ``` 64 | 65 | 66 | ### Use the MCP Server with Claude Desktop 67 | 68 | Watch the following demo video or read on for detailed instructions. 69 | 70 | 71 | 72 | https://github.com/user-attachments/assets/d45da132-46f0-48ac-a1b9-3b1b1b8fd638 73 | 74 | 75 | 76 | 1. In the Claude Desktop app, navigate to the “Settings” pane, select the “Developer” tab and click on “Edit Config”. 77 | 2. Open the `claude_desktop_config.json` file and add the following configuration to the "mcpServers" section to configure the Azure Database for PostgreSQL MCP server: 78 | 79 | ```json 80 | { 81 | "mcpServers": { 82 | "azure-postgresql-mcp": { 83 | "command": "\\azure-postgresql-mcp-venv\\Scripts\\python", 84 | "args": [ 85 | "\\azure_postgresql_mcp.py" 86 | ], 87 | "env": { 88 | "PGHOST": "", 89 | "PGUSER": "", 90 | "PGPASSWORD": "", 91 | "PGDATABASE": "" 92 | } 93 | } 94 | } 95 | } 96 | ``` 97 | **Note**: Here, we use password-based authentication to connect the MCP Server to Azure Database for PostgreSQL for testing purposes only. However, we recommend using Microsoft Entra authentication. Please refer to [these instructions](#using-microsoft-entra-authentication-method) for guidance. 98 | 3. Restart the Claude Desktop app. 99 | 4. Upon restarting, you should see a hammer icon at the bottom of the input box. Selecting this icon will display the tools provided by the MCP Server. 100 | 101 | You are now all set to start interacting with your data using natural language queries through Claude Desktop! 102 | 103 | ### Use the MCP Server with Visual Studio Code 104 | 105 | Watch the following demo video or read on for detailed instructions. 106 | 107 | 108 | 109 | https://github.com/user-attachments/assets/12328e84-7045-4e3c-beab-4936d7a20c21 110 | 111 | 112 | 113 | 1. In Visual Studio Code, navigate to “File”, select “Preferences” and then choose “Settings”. 114 | 2. Search for “MCP” and select “Edit in settings.json”. 115 | 3. Add the following configuration to the “mcp” section of the `settings.json` file: 116 | 117 | ```JSON 118 | { 119 | "mcp": { 120 | "inputs": [], 121 | "servers": { 122 | "azure-postgresql-mcp": { 123 | "command": "\\azure-postgresql-mcp-venv\\Scripts\\python", 124 | "args": [ 125 | "\\azure_postgresql_mcp.py" 126 | ], 127 | "env": { 128 | "PGHOST": "", 129 | "PGUSER": "", 130 | "PGPASSWORD": "", 131 | "PGDATABASE": "" 132 | } 133 | } 134 | } 135 | } 136 | } 137 | ``` 138 | **Note**: Here, we use password-based authentication to connect the MCP Server to Azure Database for PostgreSQL for testing purposes only. However, we recommend using Microsoft Entra authentication. Please refer to [these instructions](#using-microsoft-entra-authentication-method) for guidance. 139 | 4. Select the “Copilot” status icon in the upper-right corner to open the GitHub Copilot Chat window. 140 | 5. Choose “Agent mode” from the dropdown at the bottom of the chat input box. 141 | 5. Click on “Select Tools” (hammer icon) to view the Tools exposed by the MCP Server. 142 | 143 | You are now all set to start interacting with your data using natural language queries through VS Code! 144 | 145 | ## Using Microsoft Entra authentication method 146 | 147 | To Microsoft Entra authentication method (recommended) to connect your MCP Server to Azure Database for PostgreSQL, update the MCP Server configuration in `claude_desktop_config.json` file \(Claude Desktop\) and `settings.json` \(Visual Studio Code\) with the following code: 148 | 149 | ```json 150 | "azure-postgresql-mcp": { 151 | "command": "\\azure-postgresql-mcp-venv\\Scripts\\python", 152 | "args": [ 153 | "\\azure_postgresql_mcp.py" 154 | ], 155 | "env": { 156 | "PGHOST": "", 157 | "PGUSER": "", 158 | "AZURE_USE_AAD": "True", 159 | "AZURE_SUBSCRIPTION_ID": "", 160 | "AZURE_RESOURCE_GROUP": "" 161 | } 162 | } 163 | ``` 164 | 165 | ## Contributing 166 | The Azure Database for PostgreSQL MCP Server is currently in Preview. As we continue to develop and enhance its features, we welcome all contributions! For more details, see the [CONTRIBUTING.md](CONTRIBUTING.md) file. 167 | 168 | ## License 169 | This project is licensed under the MIT License. For more details, see the [LICENSE](LICENSE.md) file. 170 | -------------------------------------------------------------------------------- /tests/test_azure_postgresql_mcp.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import unittest 3 | from unittest.mock import MagicMock, patch 4 | 5 | # Constants 6 | NETWORK_ERROR_MESSAGE = "Network error" 7 | 8 | from azure_postgresql_mcp import AzurePostgreSQLMCP 9 | 10 | 11 | class TestAzurePostgreSQLMCPAADEnabled(unittest.TestCase): 12 | """Tests for AzurePostgreSQLMCP with AAD enabled.""" 13 | 14 | @patch("azure_postgresql_mcp.DefaultAzureCredential") 15 | @patch("azure_postgresql_mcp.PostgreSQLManagementClient") 16 | def setUp(self, mock_postgresql_client, mock_credential): 17 | # Mock the credential and client 18 | mock_credential.return_value = MagicMock() 19 | mock_client_instance = MagicMock() 20 | mock_postgresql_client.return_value = mock_client_instance 21 | 22 | """Set up the AzurePostgreSQLMCP instance with AAD enabled.""" 23 | with patch.dict( 24 | "os.environ", 25 | { 26 | "AZURE_USE_AAD": "True", 27 | "PGHOST": "test-host", 28 | "PGUSER": "test-user", 29 | "PGPASSWORD": "test-password", 30 | "AZURE_SUBSCRIPTION_ID": "test-subscription-id", 31 | "AZURE_RESOURCE_GROUP": "test-resource-group", 32 | }, 33 | ): 34 | self.azure_pg_mcp = AzurePostgreSQLMCP() 35 | self.azure_pg_mcp.init() 36 | 37 | def test_get_server_config(self): 38 | mock_server = MagicMock() 39 | mock_server.name = "test-server" 40 | mock_server.location = "eastus" 41 | mock_server.version = "12" 42 | mock_server.sku.name = "Standard_D2s_v3" 43 | mock_server.storage.storage_size_gb = 100 44 | mock_server.backup.backup_retention_days = 7 45 | mock_server.backup.geo_redundant_backup = "Enabled" 46 | 47 | # Ensure the mocked server response is serializable 48 | self.azure_pg_mcp.postgresql_client.servers.get.return_value = mock_server 49 | # Call the method 50 | result = self.azure_pg_mcp.get_server_config() 51 | 52 | # Assert the result 53 | self.assertIn("test-server", result) 54 | self.assertIn("eastus", result) 55 | self.assertIn("12", result) 56 | self.assertIn("Standard_D2s_v3", result) 57 | self.assertIn("100", result) 58 | self.assertIn("7", result) 59 | self.assertIn("Enabled", result) 60 | 61 | def test_get_server_parameter(self): 62 | # Mock the configuration response 63 | mock_configuration = MagicMock() 64 | mock_configuration.name = "max_connections" 65 | mock_configuration.value = "100" 66 | 67 | self.azure_pg_mcp.postgresql_client.configurations.get.return_value = ( 68 | mock_configuration 69 | ) 70 | 71 | # Call the method 72 | result = self.azure_pg_mcp.get_server_parameter("max_connections") 73 | 74 | # Assert the result 75 | self.assertIn("max_connections", result) 76 | self.assertIn("100", result) 77 | 78 | 79 | class TestAzurePostgreSQLMCPAADDisabled(unittest.TestCase): 80 | """Tests for AzurePostgreSQLMCP with AAD disabled.""" 81 | 82 | def setUp(self): 83 | patcher = patch.dict( 84 | "os.environ", 85 | { 86 | "PGHOST": "test-host", 87 | "PGUSER": "test-user", 88 | "PGPASSWORD": "test-password", 89 | }, 90 | ) 91 | self.addCleanup(patcher.stop) 92 | patcher.start() 93 | self.azure_pg_mcp = AzurePostgreSQLMCP() 94 | self.azure_pg_mcp.init() 95 | 96 | @patch("psycopg.connect") 97 | def test_query_data(self, mock_connect): 98 | # Mock the cursor and its behavior 99 | mock_cursor = MagicMock() 100 | mock_cursor.description = [("col1",), ("col2",)] 101 | mock_cursor.fetchall.return_value = [(1, "value1"), (2, "value2")] 102 | 103 | mock_connect.return_value.__enter__.return_value.cursor.return_value.__enter__.return_value = ( 104 | mock_cursor 105 | ) 106 | 107 | # Call the method 108 | result = self.azure_pg_mcp.query_data("test_db", "SELECT * FROM test_table;") 109 | 110 | # Assert the result 111 | self.assertIn("value1", result) 112 | self.assertIn("value2", result) 113 | 114 | @patch("psycopg.connect") 115 | def test_create_table(self, mock_connect): 116 | # Mock the connection and cursor 117 | mock_cursor = MagicMock() 118 | mock_connect.return_value.__enter__.return_value.cursor.return_value.__enter__.return_value = ( 119 | mock_cursor 120 | ) 121 | 122 | # Call the method 123 | self.azure_pg_mcp.create_table("test_db", "CREATE TABLE test_table (id INT);") 124 | 125 | # Assert that the query was executed and committed 126 | mock_cursor.execute.assert_called_once_with("CREATE TABLE test_table (id INT);") 127 | mock_connect.return_value.__enter__.return_value.commit.assert_called_once() 128 | 129 | @patch("psycopg.connect") 130 | def test_drop_table(self, mock_connect): 131 | # Mock the connection and cursor 132 | mock_cursor = MagicMock() 133 | mock_connect.return_value.__enter__.return_value.cursor.return_value.__enter__.return_value = ( 134 | mock_cursor 135 | ) 136 | 137 | # Call the method 138 | self.azure_pg_mcp.drop_table("test_db", "DROP TABLE test_table;") 139 | 140 | # Assert that the query was executed and committed 141 | mock_cursor.execute.assert_called_once_with("DROP TABLE test_table;") 142 | mock_connect.return_value.__enter__.return_value.commit.assert_called_once() 143 | 144 | 145 | class TestAzurePostgreSQLMCPNetworkErrors(unittest.TestCase): 146 | """Tests for handling network errors in AzurePostgreSQLMCP.""" 147 | 148 | @patch("azure_postgresql_mcp.DefaultAzureCredential") 149 | @patch("azure_postgresql_mcp.PostgreSQLManagementClient") 150 | def setUp(self, mock_postgresql_client, mock_credential): 151 | # Mock the credential and client 152 | mock_credential.return_value = MagicMock() 153 | mock_client_instance = MagicMock() 154 | mock_postgresql_client.return_value = mock_client_instance 155 | 156 | with patch.dict( 157 | "os.environ", 158 | { 159 | "PGHOST": "test-host", 160 | "PGUSER": "test-user", 161 | "PGPASSWORD": "test-password", 162 | "AZURE_SUBSCRIPTION_ID": "test-subscription-id", 163 | "AZURE_RESOURCE_GROUP": "test-resource-group", 164 | "AZURE_USE_AAD": "True", 165 | }, 166 | ): 167 | self.azure_pg_mcp = AzurePostgreSQLMCP() 168 | self.azure_pg_mcp.init() 169 | 170 | @patch("psycopg.connect") 171 | def test_query_data_network_error(self, mock_connect): 172 | # Simulate a network error 173 | mock_connect.side_effect = Exception(NETWORK_ERROR_MESSAGE) 174 | 175 | # Call the method 176 | result = self.azure_pg_mcp.query_data("test_db", "SELECT * FROM test_table;") 177 | 178 | # Assert the result 179 | self.assertEqual(result, "") 180 | 181 | @patch("psycopg.connect") 182 | def test_create_table_network_error(self, mock_connect): 183 | # Simulate a network error 184 | mock_connect.side_effect = Exception("Network error") 185 | 186 | # Call the method 187 | self.azure_pg_mcp.create_table("test_db", "CREATE TABLE test_table (id INT);") 188 | 189 | # Assert that no exception was raised 190 | mock_connect.return_value.__enter__.return_value.commit.assert_not_called() 191 | 192 | def test_get_server_config_network_error(self): 193 | # Simulate a network error 194 | self.azure_pg_mcp.postgresql_client.servers.get.side_effect = Exception( 195 | NETWORK_ERROR_MESSAGE 196 | ) 197 | 198 | with self.assertRaises(Exception) as context: 199 | self.azure_pg_mcp.get_server_config() 200 | 201 | # Assert the exception message 202 | self.assertEqual(str(context.exception), "Network error") 203 | 204 | def test_get_server_parameter_network_error(self): 205 | # Simulate a network error 206 | self.azure_pg_mcp.postgresql_client.configurations.get.side_effect = Exception( 207 | NETWORK_ERROR_MESSAGE 208 | ) 209 | 210 | with self.assertRaises(Exception) as context: 211 | self.azure_pg_mcp.get_server_parameter("max_connections") 212 | 213 | # Assert the exception message 214 | self.assertEqual(str(context.exception), "Network error") 215 | 216 | 217 | if __name__ == "__main__": 218 | unittest.main() 219 | -------------------------------------------------------------------------------- /src/azure_postgresql_mcp.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) Microsoft Corporation. 3 | Licensed under the MIT License. 4 | """ 5 | 6 | """ 7 | MCP server for Azure Database for PostgreSQL - Flexible Server. 8 | 9 | This server exposes the following capabilities: 10 | 11 | Tools: 12 | - create_table: Creates a table in a database. 13 | - drop_table: Drops a table in a database. 14 | - get_databases: Gets the list of all the databases in a server instance. 15 | - get_schemas: Gets schemas of all the tables. 16 | - get_server_config: Gets the configuration of a server instance. [Available with Microsoft EntraID] 17 | - get_server_parameter: Gets the value of a server parameter. [Available with Microsoft EntraID] 18 | - query_data: Runs read queries on a database. 19 | - update_values: Updates or inserts values into a table. 20 | 21 | Resources: 22 | - databases: Gets the list of all databases in a server instance. 23 | 24 | To run the code using PowerShell, expose the following variables: 25 | 26 | ``` 27 | $env:PGHOST="" 28 | $env:PGUSER="" 29 | $env:PGPASSWORD="" 30 | ``` 31 | 32 | Run the MCP Server using the following command: 33 | 34 | ``` 35 | python azure_postgresql_mcp.py 36 | ``` 37 | 38 | For detailed usage instructions, please refer to the README.md file. 39 | 40 | """ 41 | 42 | import json 43 | import logging 44 | import os 45 | import sys 46 | import urllib.parse 47 | 48 | import psycopg 49 | from azure.identity import DefaultAzureCredential 50 | from azure.mgmt.postgresqlflexibleservers import PostgreSQLManagementClient 51 | from mcp.server.fastmcp import FastMCP 52 | from mcp.server.fastmcp.resources import FunctionResource 53 | 54 | logger = logging.getLogger("azure") 55 | logger.setLevel(logging.ERROR) 56 | 57 | 58 | class AzurePostgreSQLMCP: 59 | def init(self): 60 | self.aad_in_use = os.environ.get("AZURE_USE_AAD") 61 | self.dbhost = self.get_environ_variable("PGHOST") 62 | self.dbuser = urllib.parse.quote(self.get_environ_variable("PGUSER")) 63 | 64 | if self.aad_in_use == "True": 65 | self.subscription_id = self.get_environ_variable("AZURE_SUBSCRIPTION_ID") 66 | self.resource_group_name = self.get_environ_variable("AZURE_RESOURCE_GROUP") 67 | self.server_name = ( 68 | self.dbhost.split(".", 1)[0] if "." in self.dbhost else self.dbhost 69 | ) 70 | self.credential = DefaultAzureCredential() 71 | self.postgresql_client = PostgreSQLManagementClient( 72 | self.credential, self.subscription_id 73 | ) 74 | # Password initialization should be done after checking if AAD is in use 75 | # because then we need to get the token using the credential 76 | # which is only available after the above block. 77 | self.password = self.get_password() 78 | 79 | @staticmethod 80 | def get_environ_variable(name: str): 81 | """Helper function to get environment variable or raise an error.""" 82 | value = os.environ.get(name) 83 | if value is None: 84 | raise EnvironmentError(f"Environment variable {name} not found.") 85 | return value 86 | 87 | def get_password(self) -> str: 88 | """Get password based on the auth mode set""" 89 | if self.aad_in_use == "True": 90 | return self.credential.get_token( 91 | "https://ossrdbms-aad.database.windows.net/.default" 92 | ).token 93 | else: 94 | return self.get_environ_variable("PGPASSWORD") 95 | 96 | def get_dbs_resource_uri(self): 97 | """Gets the resource URI exposed as MCP resource for getting list of dbs.""" 98 | dbhost_normalized = ( 99 | self.dbhost.split(".", 1)[0] if "." in self.dbhost else self.dbhost 100 | ) 101 | return f"flexpg://{dbhost_normalized}/databases" 102 | 103 | def get_databases_internal(self) -> str: 104 | """Internal function which gets the list of all databases in a server instance.""" 105 | try: 106 | with psycopg.connect( 107 | f"host={self.dbhost} user={self.dbuser} dbname='postgres' password={self.password}" 108 | ) as conn: 109 | with conn.cursor() as cur: 110 | cur.execute( 111 | "SELECT datname FROM pg_database WHERE datistemplate = false;" 112 | ) 113 | colnames = [desc[0] for desc in cur.description] 114 | dbs = cur.fetchall() 115 | return json.dumps( 116 | { 117 | "columns": str(colnames), 118 | "rows": "".join(str(row) for row in dbs), 119 | } 120 | ) 121 | except Exception as e: 122 | logger.error(f"Error: {str(e)}") 123 | return "" 124 | 125 | def get_databases_resource(self): 126 | """Gets list of databases as a resource""" 127 | return self.get_databases_internal() 128 | 129 | def get_databases(self): 130 | """Gets the list of all the databases in a server instance.""" 131 | return self.get_databases_internal() 132 | 133 | def get_connection_uri(self, dbname: str) -> str: 134 | """Construct URI for connection.""" 135 | return f"host={self.dbhost} dbname={dbname} user={self.dbuser} password={self.password}" 136 | 137 | def get_schemas(self, database: str): 138 | """Gets schemas of all the tables.""" 139 | try: 140 | with psycopg.connect(self.get_connection_uri(database)) as conn: 141 | with conn.cursor() as cur: 142 | cur.execute( 143 | "SELECT table_name, column_name, data_type FROM information_schema.columns " 144 | "WHERE table_schema = 'public' ORDER BY table_name, ordinal_position;" 145 | ) 146 | colnames = [desc[0] for desc in cur.description] 147 | tables = cur.fetchall() 148 | return json.dumps( 149 | { 150 | "columns": str(colnames), 151 | "rows": "".join(str(row) for row in tables), 152 | } 153 | ) 154 | except Exception as e: 155 | logger.error(f"Error: {str(e)}") 156 | return "" 157 | 158 | def query_data(self, dbname: str, s: str) -> str: 159 | """Runs read queries on a database.""" 160 | try: 161 | with psycopg.connect(self.get_connection_uri(dbname)) as conn: 162 | with conn.cursor() as cur: 163 | cur.execute(s) 164 | rows = cur.fetchall() 165 | colnames = [desc[0] for desc in cur.description] 166 | return json.dumps( 167 | { 168 | "columns": str(colnames), 169 | "rows": ",".join(str(row) for row in rows), 170 | } 171 | ) 172 | except Exception as e: 173 | logger.error(f"Error: {str(e)}") 174 | return "" 175 | 176 | def exec_and_commit(self, dbname: str, s: str) -> None: 177 | """Internal function to execute and commit transaction.""" 178 | try: 179 | with psycopg.connect(self.get_connection_uri(dbname)) as conn: 180 | with conn.cursor() as cur: 181 | cur.execute(s) 182 | conn.commit() 183 | except Exception as e: 184 | logger.error(f"Error: {str(e)}") 185 | 186 | def update_values(self, dbname: str, s: str): 187 | """Updates or inserts values into a table.""" 188 | self.exec_and_commit(dbname, s) 189 | 190 | def create_table(self, dbname: str, s: str): 191 | """Creates a table in a database.""" 192 | self.exec_and_commit(dbname, s) 193 | 194 | def drop_table(self, dbname: str, s: str): 195 | """Drops a table in a database.""" 196 | self.exec_and_commit(dbname, s) 197 | 198 | def get_server_config(self) -> str: 199 | """Gets the configuration of a server instance. [Available with Microsoft EntraID]""" 200 | if self.aad_in_use: 201 | try: 202 | server = self.postgresql_client.servers.get( 203 | self.resource_group_name, self.server_name 204 | ) 205 | return json.dumps( 206 | { 207 | "server": { 208 | "name": server.name, 209 | "location": server.location, 210 | "version": server.version, 211 | "sku": server.sku.name, 212 | "storage_profile": { 213 | "storage_size_gb": server.storage.storage_size_gb, 214 | "backup_retention_days": server.backup.backup_retention_days, 215 | "geo_redundant_backup": server.backup.geo_redundant_backup, 216 | }, 217 | }, 218 | } 219 | ) 220 | except Exception as e: 221 | logger.error(f"Failed to get PostgreSQL server configuration: {e}") 222 | raise e 223 | 224 | else: 225 | raise NotImplementedError( 226 | "This tool is available only with Microsoft EntraID" 227 | ) 228 | 229 | def get_server_parameter(self, parameter_name: str) -> str: 230 | """Gets the value of a server parameter. [Available with Microsoft EntraID]""" 231 | if self.aad_in_use: 232 | try: 233 | configuration = self.postgresql_client.configurations.get( 234 | self.resource_group_name, self.server_name, parameter_name 235 | ) 236 | return json.dumps( 237 | {"param": configuration.name, "value": configuration.value} 238 | ) 239 | except Exception as e: 240 | logger.error( 241 | f"Failed to get PostgreSQL server parameter '{parameter_name}': {e}" 242 | ) 243 | raise e 244 | else: 245 | raise NotImplementedError( 246 | "This tool is available only with Microsoft EntraID" 247 | ) 248 | 249 | 250 | if __name__ == "__main__": 251 | mcp = FastMCP("Flex PG Explorer") 252 | azure_pg_mcp = AzurePostgreSQLMCP() 253 | azure_pg_mcp.init() 254 | mcp.add_tool(azure_pg_mcp.get_databases) 255 | mcp.add_tool(azure_pg_mcp.get_schemas) 256 | mcp.add_tool(azure_pg_mcp.query_data) 257 | mcp.add_tool(azure_pg_mcp.update_values) 258 | mcp.add_tool(azure_pg_mcp.create_table) 259 | mcp.add_tool(azure_pg_mcp.drop_table) 260 | mcp.add_tool(azure_pg_mcp.get_server_config) 261 | mcp.add_tool(azure_pg_mcp.get_server_parameter) 262 | databases_resource = FunctionResource( 263 | name=azure_pg_mcp.get_dbs_resource_uri(), 264 | uri=azure_pg_mcp.get_dbs_resource_uri(), 265 | description="List of databases in the server", 266 | mime_type="application/json", 267 | fn=azure_pg_mcp.get_databases_resource, 268 | ) 269 | 270 | # Add the resource to the MCP server 271 | mcp.add_resource(databases_resource) 272 | mcp.run() 273 | --------------------------------------------------------------------------------