phab.py 7.94 KB
Newer Older
1
from __future__ import annotations
2

3
4
import json
import os
5
from collections import UserDict, UserList, deque
20after4's avatar
20after4 committed
6
7
8
9
10
from collections.abc import (
    Iterable,
    Collection,
    Mapping,
    MutableMapping,
11
12
    MutableSequence,
)
13
14
from pprint import pprint
from tokenize import Number
15
16
from typing import Sequence, Union
from operator import itemgetter
17

18
19
20
# todo: remove dependency on requests
import requests

21
from ddd.data import Data, DataIterator, wrapitem
20after4's avatar
20after4 committed
22
from ddd.phobjects import PHID, PHObject, PhabObjectBase, isPHID, json_object_hook
23

24
25
26
27
28
29
class Cursor(object):
    args: MutableMapping
    cursor: MutableMapping
    conduit: Conduit
    data: deque[Data]
    result: MutableMapping
30
    method: str
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

    def asdict(self, key="id"):
        return {obj[key]: obj for obj in self.data}

    def next_page(self):
        """
        Load the next page of results from conduit, using the cursor that was
        returned by the most recently fetched page to specify the starting
        point. This is specified by an "after" argument added to the request.
        """
        after = self.cursor.get("after", None)
        if after is None:
            raise ConduitException(
                conduit=self.conduit, message="Cannot fetch pages beyond the last."
            )
        args = self.args
        args["after"] = after
        res = self.conduit.raw_request(method=self.method, args=args)
        self.handle_result(res)

    def has_more(self):
        return self.cursor.get("after", None)

    def fetch_all(self):
        """ Sequentially Fetch all pages of results from the server. """
        while self.has_more():
            self.next_page()
        return self.data

    def handle_result(self, response):
        raise ConduitException('Not Implemented')
62
63

class Conduit(object):
64
65
    phab_url: str = "https://phabricator.wikimedia.org/api/"
    token: str = ""
66

67
    def __init__(self, phab_url: str = None, token: str = None):
68
69
70
        if phab_url:
            self.phab_url = phab_url

71
72
73
74
75
76
        if token:
            self.token = token
        else:
            self.token = self._get_token()

        if self.token is None:
77
            err = "Unable to find a conduit token in ~/.arcrc or environment"
78
            raise ConduitException(conduit=self, message=err)
79
80
81
82
83
84

    def _get_token(self):
        """
        Use the $CONDUIT_TOKEN envvar, fallback to whatever is in ~/.arcrc
        """
        token = None
85
        token_path = os.path.expanduser("~/.arcrc")
86
87
88
        if os.path.exists(token_path):
            with open(token_path) as f:
                arcrc = json.load(f)
89
90
                if self.phab_url in arcrc["hosts"]:
                    token = arcrc["hosts"][self.phab_url]["token"]
91

92
        return os.environ.get("CONDUIT_TOKEN", token)
93

20after4's avatar
20after4 committed
94
    def raw_request(self, method: str, args: Mapping) -> requests.Response:
95
        """
96
        Helper method to call a phabricator api and return a ConduitCursor
97
98
        which can be used to iterate over all of the resulting records.
        """
99
100
101
102
103
        req = flatten_for_post(args)
        req["api.token"] = self.token
        r = requests.post(f"{self.phab_url}{method}", data=req)
        return r

104
    def request(self, method: str, args: MutableMapping) -> Cursor:
105
106
107
108
        r = self.raw_request(method=method, args=args)
        return ConduitCursor(conduit=self, res=r, method=method, args=args)

    def edit(self, method: str, objectidentifier: str, transactions: list):
109
110
111
112
113
114
115
116
        """
        Calls a conduit "edit" method which applies a list of transactions
        to a specified object (specified by a phabricator identifier such as
        T123, #project or a PHID of the appropriate type)

        Raises an exception if something goes wrong and returns the decoded
        conduit response otherwise.
        """
117
        req = {
118
119
            "parameters": {
                "transactions": transactions,
120
                "objectidentifier": objectidentifier,
121
122
            }
        }
123
124
125
        req = flatten_for_post(req)
        req["api.token"] = self.token
        r = requests.post(f"{self.phab_url}{method}", data=req)
20after4's avatar
20after4 committed
126
        json = r.json(object_hook=json_object_hook)
127
128
        if json["error_info"] is not None:
            raise ConduitException(conduit=self, message=json["error_info"])
129
130
        return json

131
132
    def project_search(self, queryKey="all",
                       constraints: MutableMapping={}) -> Cursor:
133

134
135
136
137
138
139
        return self.request('project.search', {
            "queryKey": queryKey,
            "constraints": constraints
        })

class ConduitCursor(Cursor):
140
    """
141
142
143
    ConduitCursor handles fetching multiple pages of records from the conduit
    api so that one api call can be treated as a single collection of results even
    though it's split across multiple requests to the server.
144
    """
145
146
147
148
149
150
151
    def __init__(
        self,
        conduit: Conduit,
        res: requests.Response,
        method: str,
        args: MutableMapping,
    ):
152
153
154
        self.conduit = conduit
        self.method = method
        self.args = args
155
        self.cursor = {}
156
        self.data = deque()
157
158
159
        self.handle_result(res)

    def retry(self):
20after4's avatar
20after4 committed
160
161
        res = self.conduit.raw_request(method=self.method, args=self.args)
        self.handle_result(res)
162

163
    def handle_result(self, res: requests.Response):
164
165
166
167
168
        """
        Process the result from a conduit call and store the records, along
        with a cursor for fetching further pages when the result exceeds the
        limit for a single request. The default and maximum limit is 100.
        """
20after4's avatar
20after4 committed
169
        json = res.json(object_hook=json_object_hook)
170
171
        if json["error_info"] is not None:
            raise ConduitException(conduit=self.conduit, message=json["error_info"])
172

173
        self.result = json["result"]
174
175

        if "cursor" in self.result:
176
            self.cursor = self.result["cursor"]
177
178
        else:
            self.cursor = {}
179

180
        if "data" in self.result:
181
182
183
184
185
            # Modern conduit methods return a result map with the key "data"
            # mapped to a list of records and the key "cursor" maps to a record
            # of pagination details.
            self.data.extend(self.result["data"])
        else:
186
            # Older methods just return a result:
187
            self.data.extend(self.result.values())
188
189
190
191

    def __iter__(self):
        return DataIterator(self.data)

192
    def __getitem__(self, key):
20after4's avatar
20after4 committed
193
        return self.data[key]
194
195
196
197
198
199
200
201
202

    def __len__(self):
        return len(self.data)

    def __contains__(self, item):
        return item in self.data


class ConduitException(Exception):
203
204
205
    def __init__(
        self, message: str, conduit: Conduit = None, result: ConduitCursor = None
    ):
206
207
208
209
        self.conduit = conduit
        self.result = result
        self.message = message

210
211
212
213
214
    def __repr__(self):
        return "ConduitException(message='%s')" % self.message

    def __str__(self):
        return "ConduitException: " + self.message
215

20after4's avatar
20after4 committed
216

217
def flatten_for_post(h, result: dict = None, kk=None) -> dict[str, str]:
218
219
220
221
222
223
224
225
226
227
228
229
    """
    Since phab expects x-url-encoded form post data (meaning each
    individual list element is named). AND because, evidently, requests
    can't do this for me, I found a solution via stackoverflow.

    See also:
    <https://secure.phabricator.com/T12447>
    <https://stackoverflow.com/questions/26266664/requests-form-urlencoded-data/36411923>
    """
    if result is None:
        result = {}

230
    if isinstance(h, (str, bool)):
231
        result[kk] = h
232
233
234
235
236
    elif isinstance(h, (int, float)):
        # Because conduit response comes back empty when you pass raw int
        # values:
        result[kk] = str(h)
    elif hasattr(h, "items"):
237
238
        for (k, v) in h.items():
            key = k if kk is None else "%s[%s]" % (kk, k)
239
            if hasattr(v, "items"):
240
                for i, v1 in v.items():
241
                    flatten_for_post(v1, result, "%s[%s]" % (key, i))
242
243
            else:
                flatten_for_post(v, result, key)
244
245
246
    elif isinstance(h, Iterable):
        for i, v1 in enumerate(h):
            flatten_for_post(v1, result, "%s[%d]" % (kk, i))
247
    return result