Coverage for mindsdb / api / mysql / mysql_proxy / data_types / mysql_packets / binary_resultset_row_package.py: 8%

1""" 

2******************************************************* 

3 * Copyright (C) 2017 MindsDB Inc. <copyright@mindsdb.com> 

4 * 

5 * This file is part of MindsDB Server. 

6 * 

7 * MindsDB Server can not be copied and/or distributed without the express 

8 * permission of MindsDB Inc 

9 ******************************************************* 

10""" 

11 

12import datetime as dt 

13import struct 

14 

15import pandas as pd 

16 

17from mindsdb.api.mysql.mysql_proxy.data_types.mysql_datum import Datum 

18from mindsdb.api.mysql.mysql_proxy.data_types.mysql_packet import Packet 

19from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import TYPES 

20 

21 

22class BinaryResultsetRowPacket(Packet): 

23 """ 

24 Implementation based on: 

25 https://mariadb.com/kb/en/resultset-row/#binary-resultset-row 

26 https://dev.mysql.com/doc/internals/en/null-bitmap.html 

27 """ 

28 

29 def setup(self): 

30 data = self._kwargs.get("data", {}) 

31 columns = self._kwargs.get("columns", {}) 

32 

33 self.value = [b"\x00"] 

34 

35 # NOTE: according to mysql's doc offset=0 only for COM_STMT_EXECUTE, mariadb's doc does't mention that 

36 # but in fact it looks like offset=2 everywhere 

37 offset = 2 

38 nulls_bitmap = bytearray((len(columns) + offset + 7) // 8) 

39 for i, el in enumerate(data): 

40 if el is not None: 

41 continue 

42 byte_index = (i + offset) // 8 

43 bit_index = (i + offset) % 8 

44 nulls_bitmap[byte_index] |= 1 << bit_index 

45 self.value.append(bytes(nulls_bitmap)) 

46 

47 for i, col in enumerate(columns): 

48 # NOTE at this moment all types sends as strings, and it works 

49 val = data[i] 

50 if val is None: 

51 continue 

52 

53 enc = None 

54 env_val = None 

55 col_type = col["type"] 

56 if col_type == TYPES.MYSQL_TYPE_DOUBLE: 

57 enc = "<d" 

58 val = float(val) 

59 elif col_type == TYPES.MYSQL_TYPE_LONGLONG: 

60 enc = "<q" 

61 val = int(val) 

62 elif col_type == TYPES.MYSQL_TYPE_LONG: 

63 enc = "<l" 

64 val = int(val) 

65 elif col_type == TYPES.MYSQL_TYPE_FLOAT: 

66 enc = "<f" 

67 val = float(val) 

68 elif col_type == TYPES.MYSQL_TYPE_YEAR: 

69 enc = "<h" 

70 val = int(float(val)) 

71 elif col_type == TYPES.MYSQL_TYPE_SHORT: 

72 enc = "<h" 

73 val = int(val) 

74 elif col_type == TYPES.MYSQL_TYPE_TINY: 

75 enc = "<B" 

76 val = int(val) 

77 elif col_type == TYPES.MYSQL_TYPE_DATE: 

78 env_val = self.encode_date(val) 

79 elif col_type == TYPES.MYSQL_TYPE_TIMESTAMP: 

80 env_val = self.encode_date(val) 

81 elif col_type == TYPES.MYSQL_TYPE_DATETIME: 

82 env_val = self.encode_date(val) 

83 elif col_type == TYPES.MYSQL_TYPE_TIME: 

84 env_val = self.encode_time(val) 

85 elif col_type == TYPES.MYSQL_TYPE_NEWDECIMAL: 

86 enc = "string" 

87 elif col_type == TYPES.MYSQL_TYPE_VECTOR: 

88 enc = "byte" 

89 elif col_type == TYPES.MYSQL_TYPE_JSON: 

90 # json have to be encoded as byte<lenenc>, but actually for json there is no differ with string<> 

91 enc = "string" 

92 else: 

93 enc = "string" 

94 

95 if enc == "": 

96 raise Exception(f"Column with type {col_type} cant be encripted") 

97 

98 if enc == "byte": 

99 self.value.append(Datum("string", val, "lenenc").toStringPacket()) 

100 elif enc == "string": 

101 if not isinstance(val, str): 

102 val = str(val) 

103 self.value.append(Datum("string", val, "lenenc").toStringPacket()) 

104 else: 

105 if env_val is None: 

106 env_val = struct.pack(enc, val) 

107 self.value.append(env_val) 

108 

109 def encode_time(self, val: dt.time | str) -> bytes: 

110 """https://mariadb.com/kb/en/resultset-row/#time-binary-encoding""" 

111 if isinstance(val, str): 

112 try: 

113 val = dt.datetime.strptime(val, "%H:%M:%S").time() 

114 except ValueError: 

115 val = dt.datetime.strptime(val, "%H:%M:%S.%f").time() 

116 if val == dt.time(0, 0, 0): 

117 return struct.pack("<B", 0) # special case for 0 time 

118 out = struct.pack("<B", 0) # positive time 

119 out += struct.pack("<L", 0) # days 

120 out += struct.pack("<B", val.hour) 

121 out += struct.pack("<B", val.minute) 

122 out += struct.pack("<B", val.second) 

123 if val.microsecond > 0: 

124 out += struct.pack("<L", val.microsecond) 

125 len_bit = struct.pack("<B", 12) 

126 else: 

127 len_bit = struct.pack("<B", 8) 

128 return len_bit + out 

129 

130 def encode_date(self, val): 

131 # date_type = None 

132 # date_value = None 

133 

134 if isinstance(val, str): 

135 forms = [ 

136 "%Y-%m-%d", 

137 "%Y-%m-%d %H:%M:%S", 

138 "%Y-%m-%d %H:%M:%S.%f", 

139 "%Y-%m-%dT%H:%M:%S", 

140 "%Y-%m-%dT%H:%M:%S.%f", 

141 ] 

142 for f in forms: 

143 try: 

144 date_value = dt.datetime.strptime(val, f) 

145 break 

146 except ValueError: 

147 date_value = None 

148 if date_value is None: 

149 raise ValueError(f"Invalid date format: {val}") 

150 date_type = "datetime" 

151 elif isinstance(val, pd.Timestamp): 

152 date_value = val 

153 date_type = "datetime" 

154 

155 out = struct.pack("<H", date_value.year) 

156 out += struct.pack("<B", date_value.month) 

157 out += struct.pack("<B", date_value.day) 

158 

159 if date_type == "datetime": 

160 out += struct.pack("<B", date_value.hour) 

161 out += struct.pack("<B", date_value.minute) 

162 out += struct.pack("<B", date_value.second) 

163 out += struct.pack("<L", date_value.microsecond) 

164 

165 len_bit = struct.pack("<B", len(out)) 

166 return len_bit + out 

167 

168 @property 

169 def body(self): 

170 string = b"".join(self.value) 

171 self.setBody(string) 

172 return self._body 

173 

174 @staticmethod 

175 def test(): 

176 import pprint 

177 

178 pprint.pprint(str(BinaryResultsetRowPacket().get_packet_string())) 

179 

180 

181if __name__ == "__main__": 181 ↛ 182line 181 didn't jump to line 182 because the condition on line 181 was never true

182 BinaryResultsetRowPacket.test()