Skip to content

Commit ac1e417

Browse files
shane-ns1shane-ns1
authored
Merge pull request #122 from cpontes-ns1/PENG-3689
PENG-3689 - Introducing support for datasets functionality
2 parents dcd7ef0 + 0c5c68e commit ac1e417

File tree

6 files changed

+531
-3
lines changed

6 files changed

+531
-3
lines changed

‎examples/datasets.py

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
#
2+
# Copyright (c) 2014 NSONE, Inc.
3+
#
4+
# License under The MIT License (MIT). See LICENSE in project root.
5+
#
6+
7+
import time
8+
9+
from ns1 import NS1
10+
11+
# NS1 will use config in ~/.nsone by default
12+
api = NS1()
13+
14+
# to specify an apikey here instead, use:
15+
16+
# from ns1 import Config
17+
# config = Config()
18+
# config.createFromAPIKey('<<CLEARTEXT API KEY>>')
19+
# api = NS1(config=config)
20+
21+
config = api.config
22+
23+
#########################
24+
# LOAD / CREATE DATASET #
25+
#########################
26+
27+
# create a dataset
28+
dt = api.datasets().create(
29+
name="my dataset",
30+
datatype={
31+
"type": "num_queries",
32+
"scope": "account",
33+
},
34+
repeat=None,
35+
timeframe={"aggregation": "monthly", "cycles": 1},
36+
export_type="csv",
37+
recipient_emails=None,
38+
)
39+
print(dt)
40+
41+
# to load an existing dataset, get a Dataset object back
42+
dt = api.datasets().retrieve(dt.get("id"))
43+
print(dt)
44+
45+
######################
46+
# DOWNLOAD REPORTS #
47+
######################
48+
49+
while True:
50+
print("waiting for report to be generated...")
51+
time.sleep(5)
52+
53+
dt = api.datasets().retrieve(dt.get("id"))
54+
reports = dt.get("reports")
55+
if reports is None:
56+
continue
57+
58+
status = reports[0].get("status")
59+
if status == "available":
60+
print("report generation completed")
61+
break
62+
63+
if status == "failed":
64+
print("failed to generate report")
65+
exit(1)
66+
67+
report = api.datasets().retrieveReport(dt.get("id"), reports[0].get("id"))
68+
file_path = "%s.%s" % (dt.get("name"), dt.get("export_type"))
69+
70+
with open(file_path, "w") as file:
71+
file.write(report)
72+
73+
print("dataset report saved to", file_path)

‎ns1/__init__.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,15 @@ def pools(self):
293293

294294
return ns1.rest.pools.Pools(self.config)
295295

296+
def datasets(self):
297+
"""
298+
Return a new raw REST interface to Datasets resources
299+
:rtype: :py:class:`ns1.rest.datasets.Datasets`
300+
"""
301+
import ns1.rest.datasets
302+
303+
return ns1.rest.datasets.Datasets(self.config)
304+
296305
# HIGH LEVEL INTERFACE
297306
def loadZone(self, zone, callback=None, errback=None):
298307
"""

‎ns1/dataset.py

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
from ns1.rest.datasets import Datasets
2+
3+
4+
class DatasetException(Exception):
5+
pass
6+
7+
8+
class Dataset(object):
9+
"""
10+
High level object representing a dataset.
11+
"""
12+
13+
def __init__(self, config):
14+
"""
15+
Create a new high level Dataset object
16+
:param ns1.config.Config config: config object
17+
"""
18+
self._rest = Datasets(config)
19+
self.config = config
20+
self.data = None
21+
22+
def __repr__(self):
23+
return "<Dataset [%s]=%s,%s,%s,%s,%s,%s>" % (
24+
self.__getitem__("id"),
25+
self.__getitem__("name"),
26+
self.__getitem__("datatype"),
27+
self.__getitem__("repeat"),
28+
self.__getitem__("timeframe"),
29+
self.__getitem__("export_type"),
30+
self.__getitem__("recipient_emails"),
31+
)
32+
33+
def __getitem__(self, item: str):
34+
if not self.data:
35+
raise DatasetException("dataset not loaded")
36+
return self.data.get(item, None)
37+
38+
def reload(self, callback=None, errback=None):
39+
"""
40+
Reload dataset data from the API.
41+
:param callback: function call back once the call has completed
42+
:param errback: function call back if the call fails
43+
"""
44+
return self.load(reload=True, callback=callback, errback=errback)
45+
46+
def load(self, id: str = None, callback=None, errback=None, reload=False):
47+
"""
48+
Load dataset data from the API.
49+
:param str id: dataset id to load
50+
:param callback: function call back once the call has completed
51+
:param bool reload: whether to reuse the instance data instead of fetching it from the server
52+
"""
53+
if not reload and self.data:
54+
return self.data
55+
if id is None and self.data:
56+
id = self.__getitem__("id")
57+
if id is None:
58+
raise DatasetException("no dataset id: did you mean to create?")
59+
60+
def success(result: dict, *args):
61+
self.data = result
62+
if callback:
63+
return callback(self)
64+
else:
65+
return self
66+
67+
return self._rest.retrieve(id, callback=success, errback=errback)
68+
69+
def loadFromDict(self, dt: dict):
70+
"""
71+
Load dataset data from a dictionary.
72+
:param dict dt: dictionary containing *at least* either an id or domain/path/target
73+
"""
74+
if "id" in dt or (
75+
"name" in dt
76+
and "datatype" in dt
77+
and "repeat" in dt
78+
and "timeframe" in dt
79+
and "export_type" in dt
80+
and "recipient_emails" in dt
81+
):
82+
self.data = dt
83+
return self
84+
else:
85+
raise DatasetException("insufficient parameters")
86+
87+
def delete(self, callback=None, errback=None):
88+
"""
89+
Delete the dataset.
90+
:param callback: function call back once the call has completed
91+
:param errback: function call back if the call fails
92+
"""
93+
id = self.__getitem__("id")
94+
return self._rest.delete(id, callback=callback, errback=errback)
95+
96+
def create(
97+
self,
98+
name: str,
99+
datatype: dict,
100+
repeat: dict,
101+
timeframe: dict,
102+
export_type: str,
103+
recipient_emails: list,
104+
callback=None,
105+
errback=None,
106+
**kwargs
107+
):
108+
"""
109+
Create a new dataset. Pass a list of keywords and their values to
110+
configure. For the list of keywords available for dataset configuration,
111+
see :attr:`ns1.rest.datasets.Datasets.PASSTHRU_FIELDS`
112+
:param str name: the name of the dataset
113+
:param str datatype: datatype settings to define the type of data to be pulled
114+
:param str repeat: repeat settings to define recurrent reports
115+
:param str timeframe: timeframe settings for the data to be pulled
116+
:param str export_type: output format of the report
117+
:param str recipient_emails: list of user emails that will receive a copy of the report
118+
:param callback: function call back once the call has completed
119+
:param errback: function call back if the call fails
120+
"""
121+
if self.data:
122+
raise DatasetException("dataset already loaded")
123+
124+
return self._rest.create(
125+
name,
126+
datatype,
127+
repeat,
128+
timeframe,
129+
export_type,
130+
recipient_emails,
131+
callback=callback,
132+
errback=errback,
133+
**kwargs
134+
)
135+
136+
def listDatasets(self, callback=None, errback=None):
137+
"""
138+
Lists all datasets currently configured.
139+
:param callback: function call back once the call has completed
140+
:param errback: function call back if the call fails
141+
:return: a list of Dataset objects
142+
"""
143+
144+
def success(result, *args):
145+
ret = []
146+
for dt in result:
147+
ret.append(Dataset(self.config).loadFromDict(dt))
148+
if callback:
149+
return callback(ret)
150+
else:
151+
return ret
152+
153+
return Datasets(self.config).list(callback=success, errback=errback)
154+
155+
def retrieveReport(
156+
self, rp_id: str, dt_id: str = None, callback=None, errback=None
157+
):
158+
"""
159+
Retrieves a generated report given a dataset id and a report id
160+
:param str rp_id: the id of the generated report to download
161+
:param str dt_id: the id of the dataset that the above report belongs to
162+
:param callback: function call back once the call has completed
163+
:param errback: function call back if the call fails
164+
:return: generated report
165+
"""
166+
167+
if dt_id is None and self.data:
168+
dt_id = self.__getitem__("id")
169+
if dt_id is None:
170+
raise DatasetException("no dataset id: did you mean to create?")
171+
172+
return Datasets(self.config).retrieveReport(
173+
dt_id, rp_id, callback=callback, errback=errback
174+
)

‎ns1/rest/datasets.py

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
from . import resource
2+
3+
4+
class Datasets(resource.BaseResource):
5+
ROOT = "datasets"
6+
7+
PASSTHRU_FIELDS = [
8+
"name",
9+
"datatype",
10+
"repeat",
11+
"timeframe",
12+
"export_type",
13+
"recipient_emails",
14+
]
15+
16+
def _buildBody(
17+
self,
18+
name: str,
19+
datatype: dict,
20+
repeat: dict,
21+
timeframe: dict,
22+
export_type: str,
23+
recipient_emails: list,
24+
**kwargs
25+
):
26+
body = {
27+
"name": name,
28+
"datatype": datatype,
29+
"repeat": repeat,
30+
"timeframe": timeframe,
31+
"export_type": export_type,
32+
"recipient_emails": recipient_emails,
33+
}
34+
self._buildStdBody(body, kwargs)
35+
return body
36+
37+
def create(
38+
self,
39+
name: str,
40+
datatype: dict,
41+
repeat: dict,
42+
timeframe: dict,
43+
export_type: str,
44+
recipient_emails: list,
45+
callback=None,
46+
errback=None,
47+
**kwargs
48+
):
49+
body = self._buildBody(
50+
name,
51+
datatype,
52+
repeat,
53+
timeframe,
54+
export_type,
55+
recipient_emails,
56+
**kwargs
57+
)
58+
return self._make_request(
59+
"PUT",
60+
"%s" % self.ROOT,
61+
body=body,
62+
callback=callback,
63+
errback=errback,
64+
)
65+
66+
def delete(self, dtId: str, callback=None, errback=None):
67+
return self._make_request(
68+
"DELETE",
69+
"%s/%s" % (self.ROOT, dtId),
70+
callback=callback,
71+
errback=errback,
72+
)
73+
74+
def list(self, callback=None, errback=None):
75+
return self._make_request(
76+
"GET",
77+
"%s" % self.ROOT,
78+
callback=callback,
79+
errback=errback,
80+
)
81+
82+
def retrieve(self, dtId: str, callback=None, errback=None):
83+
return self._make_request(
84+
"GET",
85+
"%s/%s" % (self.ROOT, dtId),
86+
callback=callback,
87+
errback=errback,
88+
)
89+
90+
def retrieveReport(
91+
self, dtId: str, rpId: str, callback=None, errback=None
92+
):
93+
return self._make_request(
94+
"GET",
95+
"%s/%s/reports/%s" % (self.ROOT, dtId, rpId),
96+
callback=callback,
97+
errback=errback,
98+
skip_json_parsing=True,
99+
)

0 commit comments

Comments
 (0)