通过Boto3在Robot测试中建立AWS实例PostgreSQL连接失败排查

AWS CLI的ssm start-session命令底层自动处理了WebSocket连接和端口转发逻辑，但直接用Boto3调用start_sessionAPI仅完成会话创建，缺少WebSocket数据流转发的关键步骤，这就是你连不上数据库的原因。以下是具体解决方法：

核心思路 #

调用Boto3的start_session后，会返回StreamUrl、SessionId和Token，需要用这些参数建立WebSocket连接，然后实现本地端口与远程PostgreSQL端口的双向数据转发。

步骤1：安装依赖 #

确保本地环境安装websocket-client库，用于处理WebSocket连接：

pip install websocket-client

步骤2：编写Robot Framework关键字（Python实现） #

创建一个Python关键字文件，实现SSM隧道的建立和端口转发：

import websocket
import threading
import socket
from botocore.exceptions import ClientError
import boto3

class SSMDBTunnel:
    def __init__(self):
        self.ssm_client = boto3.client('ssm', region_name='你的AWS区域')
        self.ws = None
        self.local_socket = None
        self.thread = None

    def start_ssm_postgres_tunnel(self, instance_id, local_port, remote_port):
        # 启动SSM会话
        try:
            response = self.ssm_client.start_session(
                Target=instance_id,
                DocumentName='AWS-StartPortForwardingSessionToRemoteHost',
                Parameters={
                    'host': ['localhost'],  # 实例上的PostgreSQL绑定地址
                    'portNumber': [str(remote_port)],
                    'localPortNumber': [str(local_port)]
                }
            )
        except ClientError as e:
            raise Exception(f"启动SSM会话失败: {e.response['Error']['Message']}")

        # 建立WebSocket连接
        stream_url = response['StreamUrl']
        token = response['TokenValue']
        self.ws = websocket.create_connection(stream_url, header={'Authorization': f"Bearer {token}"})

        # 创建本地监听socket
        self.local_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.local_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.local_socket.bind(('localhost', local_port))
        self.local_socket.listen(1)

        # 启动线程处理连接转发
        self.thread = threading.Thread(target=self._handle_port_forwarding)
        self.thread.start()

    def _handle_port_forwarding(self):
        while True:
            # 接受本地PostgreSQL客户端连接
            client_conn, addr = self.local_socket.accept()
            # 启动线程处理双向数据转发
            threading.Thread(target=self._forward_data, args=(client_conn,)).start()

    def _forward_data(self, client_conn):
        try:
            while True:
                # 从本地客户端读取数据，发送到SSM WebSocket
                data = client_conn.recv(4096)
                if not data:
                    break
                self.ws.send_binary(data)

                # 从SSM WebSocket读取数据，发送回本地客户端
                ws_data = self.ws.recv()
                if not ws_data:
                    break
                client_conn.send(ws_data)
        except Exception as e:
            pass
        finally:
            client_conn.close()

    def stop_ssm_tunnel(self):
        if self.ws:
            self.ws.close()
        if self.local_socket:
            self.local_socket.close()
        if self.thread:
            self.thread.join()

步骤3：在Robot测试中调用关键字 #

在Robot测试用例中导入并使用该关键字：

*** Settings ***
Library    SSMDBTunnel.py

*** Test Cases ***
Test PostgreSQL Connection via SSM Tunnel
    Start SSM Postgres Tunnel    instance_id=你的EC2实例ID    local_port=5432    remote_port=5432
    # 这里执行PostgreSQL连接测试（比如用DatabaseLibrary）
    Connect To Database    psycopg2    dbname=你的数据库名    user=用户名    password=密码    host=localhost    port=5432
    # 执行查询等操作
    ${result}    Query    SELECT 1;
    Should Be Equal    ${result[0][0]}    1
    # 关闭隧道
    Stop SSM Tunnel

关键配置验证 #

• 确保EC2实例已安装SSM Agent，且附加了AmazonSSMManagedInstanceCore权限策略
• IAM用户/角色拥有ssm:StartSession和ssm:TerminateSession权限
• 实例内部的PostgreSQL配置允许本地连接（pg_hba.conf中允许localhost访问）
• 本地防火墙未阻止指定的local_port

内容的提问来源于stack exchange，提问作者Valentyna Len