diff --git a/ydb/issues.py b/ydb/issues.py index a489d4e0..1f5b46c8 100644 --- a/ydb/issues.py +++ b/ydb/issues.py @@ -1,10 +1,19 @@ # -*- coding: utf-8 -*- +from __future__ import annotations + from google.protobuf import text_format import enum import queue +import typing from . import _apis +# Workaround for good IDE and universal for runtime +if typing.TYPE_CHECKING: + from _grpc.v4.protos import ydb_issue_message_pb2, ydb_operation_pb2 +else: + from ._grpc.common.protos import ydb_issue_message_pb2, ydb_operation_pb2 + _TRANSPORT_STATUSES_FIRST = 401000 _CLIENT_STATUSES_FIRST = 402000 @@ -43,10 +52,19 @@ class StatusCode(enum.IntEnum): SESSION_POOL_EMPTY = _CLIENT_STATUSES_FIRST + 40 +# TODO: convert from proto IssueMessage +class _IssueMessage: + def __init__(self, message: str, issue_code: int, severity: int, issues) -> None: + self.message = message + self.issue_code = issue_code + self.severity = severity + self.issues = issues + + class Error(Exception): status = None - def __init__(self, message, issues=None): + def __init__(self, message: str, issues: typing.Optional[typing.Iterable[_IssueMessage]] = None): super(Error, self).__init__(message) self.issues = issues self.message = message @@ -161,14 +179,14 @@ def __init__(self, message: str): super().__init__(message) -def _format_issues(issues): +def _format_issues(issues: typing.Iterable[ydb_issue_message_pb2.IssueMessage]) -> str: if not issues: return "" return " ,".join(text_format.MessageToString(issue, as_utf8=False, as_one_line=True) for issue in issues) -def _format_response(response): +def _format_response(response: ydb_operation_pb2.Operation) -> str: fmt_issues = _format_issues(response.issues) return f"{fmt_issues} (server_code: {response.status})" @@ -196,7 +214,7 @@ def _format_response(response): } -def _process_response(response_proto): +def _process_response(response_proto: ydb_operation_pb2.Operation) -> None: if response_proto.status not in _success_status_codes: exc_obj = _server_side_error_map.get(response_proto.status) raise exc_obj(_format_response(response_proto), response_proto.issues) diff --git a/ydb/types.py b/ydb/types.py index 9771e037..cf13aac0 100644 --- a/ydb/types.py +++ b/ydb/types.py @@ -1,76 +1,94 @@ # -*- coding: utf-8 -*- +from __future__ import annotations + import abc import enum import json from . import _utilities, _apis from datetime import date, datetime, timedelta +import typing import uuid import struct from google.protobuf import struct_pb2 +from . import table + + +# Workaround for good IDE and universal for runtime +if typing.TYPE_CHECKING: + from ._grpc.v4.protos import ydb_value_pb2 +else: + from ._grpc.common.protos import ydb_value_pb2 + _SECONDS_IN_DAY = 60 * 60 * 24 _EPOCH = datetime(1970, 1, 1) -def _from_date(x, table_client_settings): +def _from_date(x: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings) -> typing.Union[date, int]: if table_client_settings is not None and table_client_settings._native_date_in_result_sets: return _EPOCH.date() + timedelta(days=x.uint32_value) return x.uint32_value -def _to_date(pb, value): +def _to_date(pb: ydb_value_pb2.Value, value: typing.Union[date, int]) -> None: if isinstance(value, date): pb.uint32_value = (value - _EPOCH.date()).days else: pb.uint32_value = value -def _from_datetime_number(x, table_client_settings): +def _from_datetime_number( + x: typing.Union[float, datetime], table_client_settings: table.TableClientSettings +) -> datetime: if table_client_settings is not None and table_client_settings._native_datetime_in_result_sets: return datetime.utcfromtimestamp(x) return x -def _from_json(x, table_client_settings): +def _from_json(x: typing.Union[str, bytearray, bytes], table_client_settings: table.TableClientSettings): if table_client_settings is not None and table_client_settings._native_json_in_result_sets: return json.loads(x) return x -def _to_uuid(value_pb, table_client_settings): +def _to_uuid(value_pb: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings) -> uuid.UUID: return uuid.UUID(bytes_le=struct.pack("QQ", value_pb.low_128, value_pb.high_128)) -def _from_uuid(pb, value): +def _from_uuid(pb: ydb_value_pb2.Value, value: uuid.UUID): pb.low_128 = struct.unpack("Q", value.bytes_le[0:8])[0] pb.high_128 = struct.unpack("Q", value.bytes_le[8:16])[0] -def _from_interval(value_pb, table_client_settings): +def _from_interval( + value_pb: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings +) -> typing.Union[timedelta, int]: if table_client_settings is not None and table_client_settings._native_interval_in_result_sets: return timedelta(microseconds=value_pb.int64_value) return value_pb.int64_value -def _timedelta_to_microseconds(value): +def _timedelta_to_microseconds(value: timedelta) -> int: return (value.days * _SECONDS_IN_DAY + value.seconds) * 1000000 + value.microseconds -def _to_interval(pb, value): +def _to_interval(pb: ydb_value_pb2.Value, value: typing.Union[timedelta, int]): if isinstance(value, timedelta): pb.int64_value = _timedelta_to_microseconds(value) else: pb.int64_value = value -def _from_timestamp(value_pb, table_client_settings): +def _from_timestamp( + value_pb: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings +) -> typing.Union[datetime, int]: if table_client_settings is not None and table_client_settings._native_timestamp_in_result_sets: return _EPOCH + timedelta(microseconds=value_pb.uint64_value) return value_pb.uint64_value -def _to_timestamp(pb, value): +def _to_timestamp(pb: ydb_value_pb2.Value, value: typing.Union[datetime, int]): if isinstance(value, datetime): pb.uint64_value = _timedelta_to_microseconds(value - _EPOCH) else: @@ -129,13 +147,15 @@ class PrimitiveType(enum.Enum): DyNumber = _apis.primitive_types.DYNUMBER, "text_value" - def __init__(self, idn, proto_field, to_obj=None, from_obj=None): + def __init__( + self, idn: ydb_value_pb2.Type.PrimitiveTypeId, proto_field: typing.Optional[str], to_obj=None, from_obj=None + ): self._idn_ = idn self._to_obj = to_obj self._from_obj = from_obj self._proto_field = proto_field - def get_value(self, value_pb, table_client_settings): + def get_value(self, value_pb: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings): """ Extracts value from protocol buffer :param value_pb: A protocol buffer @@ -149,7 +169,7 @@ def get_value(self, value_pb, table_client_settings): return getattr(value_pb, self._proto_field) - def set_value(self, pb, value): + def set_value(self, pb: ydb_value_pb2.Value, value): """ Sets value in a protocol buffer :param pb: A protocol buffer @@ -176,7 +196,9 @@ def proto(self): class DataQuery(object): __slots__ = ("yql_text", "parameters_types", "name") - def __init__(self, query_id, parameters_types, name=None): + def __init__( + self, query_id: str, parameters_types: "dict[str, ydb_value_pb2.Type]", name: typing.Optional[str] = None + ): self.yql_text = query_id self.parameters_types = parameters_types self.name = _utilities.get_query_hash(self.yql_text) if name is None else name @@ -259,7 +281,7 @@ def __str__(self): class OptionalType(AbstractTypeBuilder): __slots__ = ("_repr", "_proto", "_item") - def __init__(self, optional_type): + def __init__(self, optional_type: typing.Union[AbstractTypeBuilder, PrimitiveType]): """ Represents optional type that wraps inner type :param optional_type: An instance of an inner type @@ -291,7 +313,7 @@ def __str__(self): class ListType(AbstractTypeBuilder): __slots__ = ("_repr", "_proto") - def __init__(self, list_type): + def __init__(self, list_type: typing.Union[AbstractTypeBuilder, PrimitiveType]): """ :param list_type: List item type builder """ @@ -313,7 +335,11 @@ def __str__(self): class DictType(AbstractTypeBuilder): __slots__ = ("__repr", "__proto") - def __init__(self, key_type, payload_type): + def __init__( + self, + key_type: typing.Union[AbstractTypeBuilder, PrimitiveType], + payload_type: typing.Union[AbstractTypeBuilder, PrimitiveType], + ): """ :param key_type: Key type builder :param payload_type: Payload type builder @@ -341,7 +367,7 @@ def __init__(self): self.__elements_repr = [] self.__proto = _apis.ydb_value.Type(tuple_type=_apis.ydb_value.TupleType()) - def add_element(self, element_type): + def add_element(self, element_type: typing.Union[AbstractTypeBuilder, PrimitiveType]): """ :param element_type: Adds additional element of tuple :return: self @@ -366,7 +392,7 @@ def __init__(self): self.__members_repr = [] self.__proto = _apis.ydb_value.Type(struct_type=_apis.ydb_value.StructType()) - def add_member(self, name, member_type): + def add_member(self, name: str, member_type: typing.Union[AbstractTypeBuilder, PrimitiveType]): """ :param name: :param member_type: @@ -393,7 +419,7 @@ def __init__(self): self.__columns_repr = [] self.__proto = _apis.ydb_value.Type(struct_type=_apis.ydb_value.StructType()) - def add_column(self, name, column_type): + def add_column(self, name: str, column_type: typing.Union[AbstractTypeBuilder, PrimitiveType]): """ :param name: A column name :param column_type: A column type