hoefling
hoefling

Reputation: 66261

Count requests and update request header in requests.Session

I'm writing a python client for a RESTful API that requires a special header to be passed on each request. The header has the form: X-Seq-Count: n, where n is the sequential number of the request: first made request should have the header X-Seq-Count: 1 be present, the second one should have X-Seq-Count: 2 etc.

I'm using the requests library to handle the low level HTTP calls. What would be the best approach to track the amount of requests made and inject the custom header? What I came up with is subclassing the requests.Session and overriding the Session.prepare_request method:

class CustomSession(requests.Session):

    def __init__(self):
        super().__init__()
        self.requests_count = 0

    def prepare_request(self, request):
        # increment requests counter
        self.requests_count += 1
        # update the header
        self.headers['X-Seq-Count'] = str(self.requests_count)
        return super().prepare_request(request)

Hovewer, I'm not very happy with subclassing Session. Is there a better way? I stumbled upon the event hooks in the docs, but unsure how to use them - looking at the source code, it seems that the hooks can only be applied to the response object, not the request object?

Upvotes: 1

Views: 4221

Answers (1)

georgexsh
georgexsh

Reputation: 16624

as an alternative, you can take advantage of auth mechanism of requests, you can modify the prepared Request object:

def auth_provider(req):
    global requests_count
    requests_count += 1
    req.headers['X-Seq-Count'] = requests_count
    print('requests_count:', requests_count)
    return req

requests_count = 0
s = requests.Session()
s.auth = auth_provider

s.get('https://www.example.com')
requests.get('https://www.example.com', auth=auth_provider)

output:

requests_count: 1
requests_count: 2

however subclassing Session sounds okay to me.

Upvotes: 3

Related Questions