awslabs.postgres-mcp-server 0.0.4__py3-none-any.whl → 1.0.0__py3-none-any.whl

awslabs/postgres_mcp_server/mutable_sql_detector.py

@@ -41,82 +41,53 @@ MUTATING_KEYWORDS = {
     'ANALYZE',
 }
 
-SUSPICIOUS_PATTERNS = [
-    r"(?i)'.*?--",
-    r"(?i)'.*?or\s+1=1",
-    r'(?i)\bunion\b.*\bselect\b',
-    r'(?i)\bdrop\b',
-    r'(?i)\btruncate\b',
-    r'(?i)\bgrant\b|\brevoke\b',
-    r'(?i);',
-    r"(?i)or\s+['\"]?\d+=\d+",
-]
-
 # Compile regex pattern
 MUTATING_PATTERN = re.compile(
     r'(?i)\b(' + '|'.join(re.escape(k) for k in MUTATING_KEYWORDS) + r')\b'
 )
 
-
-def remove_comments(sql: str) -> str:
-    """Remove SQL comments from the input string.
-
-    Args:
-        sql: The SQL string to process
-
-    Returns:
-        The SQL string with all comments removed
-    """
-    sql = re.sub(r'--.*?$', '', sql, flags=re.MULTILINE)
-    sql = re.sub(r'/\*.*?\*/', '', sql, flags=re.DOTALL)
-    return sql
-
-
-def remove_strings(sql: str) -> str:
-    """Remove string literals from the SQL query.
-
-    Args:
-        sql: The SQL string to process
-
-    Returns:
-        The SQL string with all string literals removed
-    """
-    # Remove single-quoted and double-quoted string literals
-    return re.sub(r"('([^']|'')*')|(\"([^\"]|\"\")*\")", '', sql)
+SUSPICIOUS_PATTERNS = [
+    r'--.*$',  # single-line comment
+    r'/\*.*?\*/',  # multi-line comment
+    r"(?i)'.*?--",  # comment injection
+    r'(?i)\bor\b\s+\d+\s*=\s*\d+',  # numeric tautology e.g. OR 1=1
+    r"(?i)\bor\b\s*'[^']+'\s*=\s*'[^']+'",  # string tautology e.g. OR '1'='1'
+    r'(?i)\bunion\b.*\bselect\b',  # UNION SELECT
+    r'(?i)\bdrop\b',  # DROP statement
+    r'(?i)\btruncate\b',  # TRUNCATE
+    r'(?i)\bgrant\b|\brevoke\b',  # GRANT or REVOKE
+    r'(?i);',  # stacked queries
+    r'(?i)\bsleep\s*\(',  # delay-based probes
+    r'(?i)\bpg_sleep\s*\(',
+    r'(?i)\bload_file\s*\(',
+    r'(?i)\binto\s+outfile\b',
+]
 
 
 def detect_mutating_keywords(sql_text: str) -> list[str]:
     """Return a list of mutating keywords found in the SQL (excluding comments)."""
-    cleaned_sql = remove_comments(sql_text)
-    cleaned_sql = remove_strings(cleaned_sql)
-    matches = MUTATING_PATTERN.findall(cleaned_sql)
+    matches = MUTATING_PATTERN.findall(sql_text)
     return list({m.upper() for m in matches})  # Deduplicated and normalized to uppercase
 
 
-def check_sql_injection_risk(parameters: list[dict] | None) -> list[dict]:
-    """Check for potential SQL injection risks in query parameters.
+def check_sql_injection_risk(sql: str) -> list[dict]:
+    """Check for potential SQL injection risks in sql query.
 
     Args:
-        parameters: List of parameter dictionaries containing name and value pairs
+        sql: query string
 
     Returns:
-        List of dictionaries containing detected security issues
+        dictionaries containing detected security issue
     """
     issues = []
-
-    if parameters is not None:
-        for param in parameters:
-            value = next(iter(param['value'].values()))
-            for pattern in SUSPICIOUS_PATTERNS:
-                if re.search(pattern, str(value)):
-                    issues.append(
-                        {
-                            'type': 'parameter',
-                            'parameter_name': param['name'],
-                            'message': f'Suspicious pattern in value: {value}',
-                            'severity': 'high',
-                        }
-                    )
-                    break
-
+    for pattern in SUSPICIOUS_PATTERNS:
+        if re.search(pattern, sql):
+            issues.append(
+                {
+                    'type': 'sql',
+                    'message': f'Suspicious pattern in query: {sql}',
+                    'severity': 'high',
+                }
+            )
+            break
     return issues

awslabs/postgres_mcp_server/server.py

@@ -29,18 +29,21 @@ from typing import Annotated, Any, Dict, List, Optional
 client_error_code_key = 'run_query ClientError code'
 unexpected_error_key = 'run_query unexpected error'
 write_query_prohibited_key = 'Your MCP tool only allows readonly query. If you want to write, change the MCP configuration per README.md'
+query_comment_prohibited_key = 'The comment in query is prohibited because of injection risk'
+query_injection_risk_key = 'Your query contains risky injection patterns'
 
 
 class DummyCtx:
     """A dummy context class for error handling in MCP tools."""
 
-    def error(self, message):
+    async def error(self, message):
         """Raise a runtime error with the given message.
 
         Args:
             message: The error message to include in the runtime error
         """
-        raise RuntimeError(f'MCP Tool Error: {message}')
+        # Do nothing
+        pass
 
 
 class DBConnection:
@@ -214,48 +217,52 @@ async def run_query(
             await ctx.error(write_query_prohibited_key)
             return [{'error': write_query_prohibited_key}]
 
-    if query_parameters is not None:
-        issues = check_sql_injection_risk(query_parameters)
-        if issues:
-            logger.info(
-                f'query is rejected because it contains risky SQL pattern, SQL query: {sql}, reasons: {issues}'
-            )
-            await ctx.error(
-                str({'message': 'Query parameter contains suspicious pattern', 'details': issues})
-            )
-            return [{'error': write_query_prohibited_key}]
+    issues = check_sql_injection_risk(sql)
+    if issues:
+        logger.info(
+            f'query is rejected because it contains risky SQL pattern, SQL query: {sql}, reasons: {issues}'
+        )
+        await ctx.error(
+            str({'message': 'Query parameter contains suspicious pattern', 'details': issues})
+        )
+        return [{'error': query_injection_risk_key}]
 
     try:
-        logger.info(f'run_query: {sql}')
-
-        execute_params = {
-            'resourceArn': db_connection.cluster_arn,
-            'secretArn': db_connection.secret_arn,
-            'database': db_connection.database,
-            'sql': sql,
-            'includeResultMetadata': True,
-        }
-
-        if query_parameters:
-            execute_params['parameters'] = query_parameters
+        logger.info(f'run_query: readonly:{db_connection.readonly_query}, SQL:{sql}')
 
-        response = await asyncio.to_thread(
-            db_connection.data_client.execute_statement, **execute_params
-        )
+        if db_connection.readonly_query:
+            response = await asyncio.to_thread(
+                execute_readonly_query, db_connection, sql, query_parameters
+            )
+        else:
+            execute_params = {
+                'resourceArn': db_connection.cluster_arn,
+                'secretArn': db_connection.secret_arn,
+                'database': db_connection.database,
+                'sql': sql,
+                'includeResultMetadata': True,
+            }
+
+            if query_parameters:
+                execute_params['parameters'] = query_parameters
+
+            response = await asyncio.to_thread(
+                db_connection.data_client.execute_statement, **execute_params
+            )
 
         logger.success('run_query successfully executed query:{}', sql)
         return parse_execute_response(response)
     except ClientError as e:
-        logger.error(f'{client_error_code_key}: {e.response["Error"]["Message"]}')
+        logger.exception(client_error_code_key)
         await ctx.error(
             str({'code': e.response['Error']['Code'], 'message': e.response['Error']['Message']})
         )
-        return [{'error': write_query_prohibited_key}]
+        return [{'error': client_error_code_key}]
     except Exception as e:
+        logger.exception(unexpected_error_key)
         error_details = f'{type(e).__name__}: {str(e)}'
-        logger.error(f'{unexpected_error_key}: {error_details}')
         await ctx.error(str({'message': error_details}))
-        return [{'error': write_query_prohibited_key}]
+        return [{'error': unexpected_error_key}]
 
 
 @mcp.tool(
@@ -276,7 +283,7 @@ async def get_table_schema(
     """
     logger.info(f'get_table_schema: {table_name}')
 
-    sql = f"""
+    sql = """
         SELECT
             a.attname AS column_name,
             pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
@@ -284,13 +291,77 @@ async def get_table_schema(
         FROM
             pg_attribute a
         WHERE
-            a.attrelid = '{table_name}'::regclass
+            a.attrelid = :table_name::regclass
             AND a.attnum > 0
             AND NOT a.attisdropped
         ORDER BY a.attnum
-    """  # nosec B608: injection risk is handled inside run_query
+    """
+
+    params = [{'name': 'table_name', 'value': {'stringValue': table_name}}]
 
-    return await run_query(sql, ctx)
+    return await run_query(sql=sql, ctx=ctx, query_parameters=params)
+
+
+def execute_readonly_query(
+    db_connection: DBConnection, query: str, parameters: Optional[List[Dict[str, Any]]] = None
+) -> dict:
+    """Execute a query under readonly transaction.
+
+    Args:
+        db_connection: connection object
+        query: query to run
+        parameters: parameters
+
+    Returns:
+        List of dictionary that contains query response rows
+    """
+    tx_id = ''
+    try:
+        # Begin read-only transaction
+        tx = db_connection.data_client.begin_transaction(
+            resourceArn=db_connection.cluster_arn,
+            secretArn=db_connection.secret_arn,
+            database=db_connection.database,
+        )
+
+        tx_id = tx['transactionId']
+
+        db_connection.data_client.execute_statement(
+            resourceArn=db_connection.cluster_arn,
+            secretArn=db_connection.secret_arn,
+            database=db_connection.database,
+            sql='SET TRANSACTION READ ONLY',
+            transactionId=tx_id,
+        )
+
+        execute_params = {
+            'resourceArn': db_connection.cluster_arn,
+            'secretArn': db_connection.secret_arn,
+            'database': db_connection.database,
+            'sql': query,
+            'includeResultMetadata': True,
+            'transactionId': tx_id,
+        }
+
+        if parameters is not None:
+            execute_params['parameters'] = parameters
+
+        result = db_connection.data_client.execute_statement(**execute_params)
+
+        db_connection.data_client.commit_transaction(
+            resourceArn=db_connection.cluster_arn,
+            secretArn=db_connection.secret_arn,
+            transactionId=tx_id,
+        )
+        return result
+    except Exception as e:
+        if tx_id:
+            db_connection.data_client.rollback_transaction(
+                resourceArn=db_connection.cluster_arn,
+                secretArn=db_connection.secret_arn,
+                transactionId=tx_id,
+            )
+        raise e
 
 
 def main():
@@ -301,8 +372,6 @@ def main():
     parser = argparse.ArgumentParser(
         description='An AWS Labs Model Context Protocol (MCP) server for postgres'
     )
-    parser.add_argument('--sse', action='store_true', help='Use SSE transport')
-    parser.add_argument('--port', type=int, default=8888, help='Port to run the server on')
     parser.add_argument('--resource_arn', required=True, help='ARN of the RDS cluster')
     parser.add_argument(
         '--secret_arn',
@@ -331,31 +400,26 @@ def main():
         DBConnectionSingleton.initialize(
             args.resource_arn, args.secret_arn, args.database, args.region, args.readonly
         )
-    except BotoCoreError as e:
-        logger.error(
-            f'Failed to RDS API client object for Postgres. Exit the MCP server. error: {str(e)}'
-        )
+    except BotoCoreError:
+        logger.exception('Failed to RDS API client object for Postgres. Exit the MCP server')
         sys.exit(1)
 
     # Test RDS API connection
     ctx = DummyCtx()
-    try:
-        asyncio.run(run_query('SELECT 1', ctx))
-    except Exception as e:
-        logger.error(
-            f'Failed to validate RDS API db connection to Postgres. Exit the MCP server. error: {e}'
-        )
+    response = asyncio.run(run_query('SELECT 1', ctx))
+    if (
+        isinstance(response, list)
+        and len(response) == 1
+        and isinstance(response[0], dict)
+        and 'error' in response[0]
+    ):
+        logger.error('Failed to validate RDS API db connection to Postgres. Exit the MCP server')
         sys.exit(1)
 
     logger.success('Successfully validated RDS API db connection to Postgres')
 
-    # Run server with appropriate transport
-    if args.sse:
-        mcp.settings.port = args.port
-        mcp.run(transport='sse')
-    else:
-        logger.info('Starting Postgres MCP server')
-        mcp.run()
+    logger.info('Starting Postgres MCP server')
+    mcp.run()
 
 
 if __name__ == '__main__':

{awslabs_postgres_mcp_server-0.0.4.dist-info → awslabs_postgres_mcp_server-1.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: awslabs.postgres-mcp-server
-Version: 0.0.4
+Version: 1.0.0
 Summary: An AWS Labs Model Context Protocol (MCP) server for postgres
 Project-URL: homepage, https://awslabs.github.io/mcp/
 Project-URL: docs, https://awslabs.github.io/mcp/servers/postgres-mcp-server/

awslabs_postgres_mcp_server-1.0.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+awslabs/__init__.py,sha256=47wJeKcStxEJwX7SVVV2pnAWYR8FxcaYoT3YTmZ5Plg,674
+awslabs/postgres_mcp_server/__init__.py,sha256=J9QAlEAe6DiFJ9irk9cOaM7NINate24hfP1RFiywKww,616
+awslabs/postgres_mcp_server/mutable_sql_detector.py,sha256=1Rywe9i7zkAT2JxG_QIc1nr1v6J49yUjfwVTvVpZe3A,2655
+awslabs/postgres_mcp_server/server.py,sha256=frD6WP6Iixdzs8qOh-2vAXkbOxVNNp_TuQYxVVcQPJU,14440
+awslabs_postgres_mcp_server-1.0.0.dist-info/METADATA,sha256=4sXbcBYaEGxKpxTf-WRuGcX3Sw1b2uJ-D9cKCt0z-w8,4897
+awslabs_postgres_mcp_server-1.0.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+awslabs_postgres_mcp_server-1.0.0.dist-info/entry_points.txt,sha256=qHCxq_MTANxTB8-mA7Wl6H71qWEMwmG-I6_koII4AXY,88
+awslabs_postgres_mcp_server-1.0.0.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
+awslabs_postgres_mcp_server-1.0.0.dist-info/licenses/NOTICE,sha256=qMIIe3h7I1Y4-CKejn50wbSKXEZLWhYHdKaRwKdXN9M,95
+awslabs_postgres_mcp_server-1.0.0.dist-info/RECORD,,

awslabs_postgres_mcp_server-0.0.4.dist-info/RECORD

@@ -1,10 +0,0 @@
-awslabs/__init__.py,sha256=47wJeKcStxEJwX7SVVV2pnAWYR8FxcaYoT3YTmZ5Plg,674
-awslabs/postgres_mcp_server/__init__.py,sha256=J9QAlEAe6DiFJ9irk9cOaM7NINate24hfP1RFiywKww,616
-awslabs/postgres_mcp_server/mutable_sql_detector.py,sha256=ciO8aCR06IZxzC6ui-brEqqou0YCO-iYX-ZFwNclQHw,3396
-awslabs/postgres_mcp_server/server.py,sha256=WSKNaoxgcnq17Bcjiom1qhA6DGN_rqlq7yYAmiKXzmU,12432
-awslabs_postgres_mcp_server-0.0.4.dist-info/METADATA,sha256=YwkdfUcb4EbxCn35HVPH-Wp0P2jSMLEBljWT7uNiv-U,4897
-awslabs_postgres_mcp_server-0.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-awslabs_postgres_mcp_server-0.0.4.dist-info/entry_points.txt,sha256=qHCxq_MTANxTB8-mA7Wl6H71qWEMwmG-I6_koII4AXY,88
-awslabs_postgres_mcp_server-0.0.4.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
-awslabs_postgres_mcp_server-0.0.4.dist-info/licenses/NOTICE,sha256=qMIIe3h7I1Y4-CKejn50wbSKXEZLWhYHdKaRwKdXN9M,95
-awslabs_postgres_mcp_server-0.0.4.dist-info/RECORD,,