Skip to content

Add Async IO support #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 78 additions & 55 deletions sumologic/async.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,94 +3,102 @@
import logging
import asyncio
import aiohttp
import threading

class SumoLogic(object):

class SumoLogic(object):
def __init__(self, accessId, accessKey, endpoint=None, cookieFile='cookies.txt'):
Copy link
Author

@dszakallas dszakallas Nov 10, 2017

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

there are a couple of rough edges here:

  • cookies have to be added back if they are really needed
  • remove explicitly setting the read timeout to infinity. I don't why but while fetching results for some search jobs I frequently got TimeOutErrors. I think this should be configurable at least.

self.session = aiohttp.ClientSession(auth=aiohttp.BasicAuth(accessId, accessKey),
read_timeout=None,
headers={'content-type': 'application/json', 'accept': 'application/json'})
self.session = aiohttp.ClientSession(
auth=aiohttp.BasicAuth(accessId, accessKey), headers={
'content-type': 'application/json',
'accept': 'application/json'
})

self.endpoint = None
self._lock = asyncio.Lock()
self._endpoint = None

async def _guard_endpoint(self):
"""
SumoLogic REST API endpoint changes based on the geo location of the client.
For example, If the client geolocation is Australia then the REST end point is
https://api.au.sumologic.com/api/v1
async def __aenter__(self):
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for safely closing the connection

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also it maybe wort exposing a close method.

await self.session.__aenter__()
return self

When the default REST endpoint (https://api.sumologic.com/api/v1) is used the server
responds with a 401 and causes the SumoLogic class instantiation to fail and this very
unhelpful message is shown 'Full authentication is required to access this resource'
async def __aexit__(self, *args):
await self.session.__aexit__(*args)

This method makes a request to the default REST endpoint and resolves the 401 to learn
the right endpoint
"""
self.endpoint = 'https://api.sumologic.com/api/v1'
response = await self.session.get('https://api.sumologic.com/api/v1/collectors') # Dummy call to get endpoint
self.endpoint = str(response.url).replace('/collectors', '') # dirty hack to sanitise URI and retain domain
async def _set_endpoint(self):
if self._endpoint is None:
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

double checked locking to avoid sending the request more than once

with await self._lock:
if self._endpoint is None:
async with self.session.get('https://api.sumologic.com/api/v1/collectors') as resp: # Dummy call to get endpoint
self._endpoint = str(resp.url).replace('/collectors', '') # dirty hack to sanitise URI and retain domain

async def delete(self, method, params=None):
await self._guard_endpoint()
r = await self.session.delete(self.endpoint + method, params=params)
await self._set_endpoint()
r = await self.session.delete(self._endpoint + method, params=params)
r.raise_for_status()
return r

async def get(self, method, params=None):
await self._guard_endpoint()
r = await self.session.get(self.endpoint + method, params=params)
await self._set_endpoint()
r = await self.session.get(self._endpoint + method, params=params)
if 400 <= r.status < 600:
r.reason = await r.text()
r.raise_for_status()
return r

async def post(self, method, params, headers=None):
await self._guard_endpoint()
r = await self.session.post(self.endpoint + method, data=json.dumps(params), headers=headers)
await self._set_endpoint()
r = await self.session.post(self._endpoint + method, data=json.dumps(params), headers=headers)
r.raise_for_status()
return r

async def put(self, method, params, headers=None):
await self._guard_endpoint()
r = await self.session.put(self.endpoint + method, data=json.dumps(params), headers=headers)
await self._set_endpoint()
r = await self.session.put(self._endpoint + method, data=json.dumps(params), headers=headers)
r.raise_for_status()
return r

async def search(self, query, fromTime=None, toTime=None, timeZone='UTC'):
params = {'q': query, 'from': fromTime, 'to': toTime, 'tz': timeZone}
r = await self.get('/logs/search', params)
return json.loads(await r.text())
async with r:
return json.loads(await r.text())

async def search_job(self, query, fromTime=None, toTime=None, timeZone='UTC'):
params = {'query': query, 'from': fromTime, 'to': toTime, 'timeZone': timeZone}
r = await self.post('/search/jobs', params)
return json.loads(await r.text())
async with r:
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

prevents leaking response objects

return json.loads(await r.text())

async def search_job_status(self, search_job):
r = await self.get('/search/jobs/' + str(search_job['id']))
return json.loads(await r.text())
async with r:
return json.loads(await r.text())

async def search_job_messages(self, search_job, limit=None, offset=0):
params = {'limit': limit, 'offset': offset}
r = await self.get('/search/jobs/' + str(search_job['id']) + '/messages', params)
return json.loads(await r.text())
async with r:
return json.loads(await r.text())

async def search_job_records(self, search_job, limit=None, offset=0):
params = {'limit': limit, 'offset': offset}
r = await self.get('/search/jobs/' + str(search_job['id']) + '/records', params)
return json.loads(await r.text())
async with r:
return json.loads(await r.text())

async def delete_search_job(self, search_job):
return await self.delete('/search/jobs/' + str(search_job['id']))

async def collectors(self, limit=None, offset=None):
params = {'limit': limit, 'offset': offset}
r = await self.get('/collectors', params)
return json.loads(await r.text())['collectors']
async with r:
return json.loads(await r.text())['collectors']

async def collector(self, collector_id):
r = await self.get('/collectors/' + str(collector_id))
return json.loads(await r.text()), r.headers['etag']
async with r:
r = await self.get('/collectors/' + str(collector_id))
return json.loads(await r.text()), r.headers['etag']

async def update_collector(self, collector, etag):
headers = {'If-Match': etag}
Expand All @@ -101,12 +109,14 @@ async def delete_collector(self, collector):

async def sources(self, collector_id, limit=None, offset=None):
params = {'limit': limit, 'offset': offset}
r = await self.get('/collectors/' + str(collector_id) + '/sources', params)
return json.loads(await r.text())['sources']
async with r:
r = await self.get('/collectors/' + str(collector_id) + '/sources', params)
return json.loads(await r.text())['sources']

async def source(self, collector_id, source_id):
r = await self.get('/collectors/' + str(collector_id) + '/sources/' + str(source_id))
return json.loads(await r.text()), r.headers['etag']
async with r:
r = await self.get('/collectors/' + str(collector_id) + '/sources/' + str(source_id))
return json.loads(await r.text()), r.headers['etag']

async def create_source(self, collector_id, source):
return await self.post('/collectors/' + str(collector_id) + '/sources', source)
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i cannot prevent leaking here without changing the API, as this returns the request object itself.

Expand All @@ -120,43 +130,56 @@ async def delete_source(self, collector_id, source):

async def create_content(self, path, data):
r = await self.post('/content/' + path, data)
return await r.text()
async with r:
return await r.text()

async def get_content(self, path):
r = await self.get('/content/' + path)
return json.loads(await r.text())
async with r:
return json.loads(await r.text())

async def delete_content(self):
r = await self.delete('/content/' + path)
return json.loads(await r.text())
async with r:
return json.loads(await r.text())

async def dashboards(self, monitors=False):
params = {'monitors': monitors}
r = await self.get('/dashboards', params)
return json.loads(await r.text())['dashboards']
async with r:
return json.loads(await r.text())['dashboards']

async def dashboard(self, dashboard_id):
r = await self.get('/dashboards/' + str(dashboard_id))
return json.loads(await r.text())['dashboard']
r = await self.get('/dashboard' + str(dashboard_id))
async with r:
return json.loads(await r.text())['dashboard']

async def dashboard_data(self, dashboard_id):
r = await self.get('/dashboards/' + str(dashboard_id) + '/data')
return json.loads(await r.text())['dashboardMonitorDatas']
async with r:
return json.loads(await r.text())['dashboardMonitorDatas']

async def search_metrics(self, query, fromTime=None, toTime=None, requestedDataPoints=600, maxDataPoints=800):
'''Perform a single Sumo metrics query'''

def millisectimestamp(ts):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe the common functions should be extracted to a util namespace to reduce code duplications.

'''Convert UNIX timestamp to milliseconds'''
if ts > 10**12:
ts = ts/(10**(len(str(ts))-13))
ts = ts / (10**(len(str(ts)) - 13))
else:
ts = ts*10**(12-len(str(ts)))
ts = ts * 10**(12 - len(str(ts)))
return int(ts)

params = {'query': [{"query":query, "rowId":"A"}],
'startTime': millisectimestamp(fromTime),
'endTime': millisectimestamp(toTime),
'requestedDataPoints': requestedDataPoints,
'maxDataPoints': maxDataPoints}
r = await self.post('/metrics/results', params)
return json.loads(await r.text())
params = {
'query': [{
"query": query,
"rowId": "A"
}],
'startTime': millisectimestamp(fromTime),
'endTime': millisectimestamp(toTime),
'requestedDataPoints': requestedDataPoints,
'maxDataPoints': maxDataPoints
}
r = await self.post('/metrics/results', params)
async with r:
return json.loads(await r.text())