ydb-python-sdk/ydb/operation.py at 6b9c807def5df44fd0048c3944a5008617fe5fd1 · ydb-platform/ydb-python-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
# -*- coding: utf-8 -*-
from . import issues
from . import _apis


def _forget_operation_request(operation_id):
    request = _apis.ydb_operation.ForgetOperationRequest(id=operation_id)
    return request


def _forget_operation_response(rpc_state, response):  # pylint: disable=W0613
    issues._process_response(response)


def _cancel_operation_request(operation_id):
    request = _apis.ydb_operation.CancelOperationRequest(id=operation_id)
    return request


def _cancel_operation_response(rpc_state, response):  # pylint: disable=W0613
    issues._process_response(response)


def _get_operation_request(self):
    request = _apis.ydb_operation.GetOperationRequest(id=self.id)
    return request


def _list_operations_response(rpc_state, response):  # pylint: disable=W0613
    issues._process_response(response)


def _list_operations_request(kind: str, page_size: int, page_token: str):
    return _apis.ydb_operation.ListOperationsRequest(kind=kind, page_size=page_size, page_token=page_token)


class OperationClient(object):
    def __init__(self, driver):
        self._driver = driver

    def cancel(self, operation_id, settings=None):
        return self._driver(
            _cancel_operation_request(operation_id),
            _apis.OperationService.Stub,
            _apis.OperationService.CancelOperation,
            _cancel_operation_response,
            settings,
        )

    def forget(self, operation_id, settings=None):
        return self._driver(
            _forget_operation_request(operation_id),
            _apis.OperationService.Stub,
            _apis.OperationService.ForgetOperation,
            _forget_operation_response,
            settings,
        )

    def list(self, kind: str, page_size: int, page_token: str, settings=None):
        return self._driver(
            _list_operations_request(kind, page_size, page_token),
            _apis.OperationService.Stub,
            _apis.OperationService.ListOperations,
            _list_operations_response,
            settings,
        )


class Operation(object):
    __slots__ = ("id", "_driver", "self_cls")

    def __init__(self, rpc_state, response, driver=None):  # pylint: disable=W0613
        # implement proper interface a bit later
        issues._process_response(response.operation)
        self.id = response.operation.id
        self._driver = driver
        # self.ready = operation.ready

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return "<Operation %s>" % (self.id,)

    def _ensure_implements(self):
        if self._driver is None:
            raise ValueError("Operation doesn't implement request!")

    def cancel(self, settings=None):
        self._ensure_implements()
        return self._driver(
            _cancel_operation_request(self.id),
            _apis.OperationService.Stub,
            _apis.OperationService.CancelOperation,
            _cancel_operation_response,
            settings,
        )

    def forget(self, settings=None):
        self._ensure_implements()
        return self._driver(
            _forget_operation_request(self.id),
            _apis.OperationService.Stub,
            _apis.OperationService.ForgetOperation,
            _forget_operation_response,
            settings,
        )

    def get(self, settings=None):
        self._ensure_implements()
        return self._driver(
            _get_operation_request(self),
            _apis.OperationService.Stub,
            _apis.OperationService.GetOperation,
            self.__class__,
            settings,
            (self._driver,),
        )

    def list(self, kind: str, page_size: int, page_token: str, settings=None):
        self._ensure_implements()
        return self._driver(
            _list_operations_request(kind, page_size, page_token),
            _apis.OperationService.Stub,
            _apis.OperationService.ListOperations,
            self.__class__,
            settings,
            (self._driver,),
        )