1# -*- coding: UTF-8 -*- 2import re 3 4from sqlalchemy import select, func, and_, or_ 5 6from pychess.Utils.const import FEN_START, WHITE, BLACK, reprResult 7from pychess.Database import model as dbmodel 8from pychess.Database.model import game, event, site, player, pl1, pl2, annotator, source, tag_game 9 10 11count_games = select([func.count()]).select_from(game) 12 13 14def parseDateTag(tag): 15 elements = re.match(r"^([0-9\?]{4})(\.([0-9\?]{2})(\.([0-9\?]{2}))?)?$", tag) 16 if elements is None: 17 y, m, d = None, None, None 18 else: 19 elements = elements.groups() 20 try: 21 y = int(elements[0]) 22 except Exception: 23 y = None 24 try: 25 m = int(elements[2]) 26 except Exception: 27 m = None 28 try: 29 d = int(elements[4]) 30 except Exception: 31 d = None 32 return y, m, d 33 34 35def save(path, model, offset, flip=False): 36 game_event = model.tags["Event"] 37 game_site = model.tags["Site"] 38 39 date = model.tags["Date"] 40 41 game_round = model.tags["Round"] 42 43 white = repr(model.players[WHITE]) 44 black = repr(model.players[BLACK]) 45 46 result = model.status 47 eco = model.tags["ECO"] 48 49 time_control = model.tags["TimeControl"] 50 board = int(model.tags["Board"]) if "Board" in model.tags else 0 51 52 white_elo = model.tags["WhiteElo"] 53 black_elo = model.tags["BlackElo"] 54 55 variant = model.variant.variant 56 57 fen = model.boards[0].board.asFen() 58 fen = fen if fen != FEN_START else "" 59 60 game_annotator = model.tags["Annotator"] if "Annotator" in model.tags else "" 61 ply_count = model.ply - model.lowply 62 63 def get_id(table, name): 64 if not name: 65 return None 66 67 selection = select([table.c.id], table.c.name == name) 68 result = conn.execute(selection) 69 id_ = result.scalar() 70 if id_ is None: 71 result = conn.execute(table.insert().values(name=name)) 72 id_ = result.inserted_primary_key[0] 73 return id_ 74 75 engine = dbmodel.get_engine(path) 76 77 conn = engine.connect() 78 trans = conn.begin() 79 try: 80 event_id = get_id(event, game_event) 81 site_id = get_id(site, game_site) 82 white_id = get_id(player, white) 83 black_id = get_id(player, black) 84 annotator_id = get_id(annotator, game_annotator) 85 86 new_values = { 87 'offset': offset, 88 'offset8': (offset >> 3) << 3, 89 'event_id': event_id, 90 'site_id': site_id, 91 'date': date, 92 'round': game_round, 93 'white_id': white_id, 94 'black_id': black_id, 95 'result': result, 96 'white_elo': white_elo, 97 'black_elo': black_elo, 98 'ply_count': ply_count, 99 'eco': eco, 100 'time_control': time_control, 101 'board': board, 102 'fen': fen, 103 'variant': variant, 104 'annotator_id': annotator_id, 105 } 106 107 if hasattr(model, "game_id") and model.game_id is not None: 108 result = conn.execute(game.update().where( 109 game.c.id == model.game_id).values(new_values)) 110 else: 111 result = conn.execute(game.insert().values(new_values)) 112 model.game_id = result.inserted_primary_key[0] 113 114 trans.commit() 115 except Exception: 116 trans.rollback() 117 raise 118 conn.close() 119 120 121col2label = {game.c.id: "Id", 122 game.c.offset: "Offset", 123 game.c.offset8: "Offset8", 124 pl1.c.name: "White", 125 pl2.c.name: "Black", 126 game.c.result: "Result", 127 event.c.name: "Event", 128 site.c.name: "Site", 129 game.c.round: "Round", 130 game.c.date: "Date", 131 game.c.white_elo: "WhiteElo", 132 game.c.black_elo: "BlackElo", 133 game.c.ply_count: "PlyCount", 134 game.c.eco: "ECO", 135 game.c.time_control: "TimeControl", 136 game.c.board: "Board", 137 game.c.fen: "FEN", 138 game.c.variant: "Variant", 139 annotator.c.name: "Annotator", 140 } 141 142 143class TagDatabase: 144 def __init__(self, engine): 145 self.engine = engine 146 147 self.cols = [col.label(col2label[col]) for col in col2label] 148 149 self.from_obj = [ 150 game.outerjoin(pl1, game.c.white_id == pl1.c.id) 151 .outerjoin(pl2, game.c.black_id == pl2.c.id) 152 .outerjoin(event, game.c.event_id == event.c.id) 153 .outerjoin(site, game.c.site_id == site.c.id) 154 .outerjoin(annotator, game.c.annotator_id == annotator.c.id)] 155 156 self.select = select(self.cols, from_obj=self.from_obj) 157 158 self.select_offsets = select([game.c.offset, ], from_obj=self.from_obj) 159 160 self.colnames = self.engine.execute(self.select).keys() 161 162 self.query = self.select 163 self.order_cols = (game.c.offset, game.c.offset) 164 self.is_desc = False 165 self.where_tags = None 166 self.where_offs = None 167 self.where_offs8 = None 168 169 def get_count(self): 170 return self.engine.execute(count_games).scalar() 171 count = property(get_count) 172 173 def close(self): 174 self.engine.dispose() 175 176 def build_order_by(self, order_col, is_desc): 177 self.is_desc = is_desc 178 self.order_cols = (order_col, game.c.offset) 179 180 def build_where_tags(self, tag_query): 181 if tag_query is not None: 182 tags = [] 183 if "white" in tag_query: 184 if "ignore_tag_colors" in tag_query: 185 tags.append(or_(pl1.c.name.like("%%%s%%" % tag_query["white"]), 186 pl2.c.name.like("%%%s%%" % tag_query["white"]))) 187 else: 188 tags.append(pl1.c.name.like("%%%s%%" % tag_query["white"])) 189 190 if "black" in tag_query: 191 if "ignore_tag_colors" in tag_query: 192 tags.append(or_(pl1.c.name.like("%%%s%%" % tag_query["black"]), 193 pl2.c.name.like("%%%s%%" % tag_query["black"]))) 194 else: 195 tags.append(pl2.c.name.like("%%%s%%" % tag_query["black"])) 196 197 if "event" in tag_query: 198 tags.append(event.c.name.like("%%%s%%" % tag_query["event"])), 199 200 if "site" in tag_query: 201 tags.append(site.c.name.like("%%%s%%" % tag_query["site"])), 202 203 if "eco_from" in tag_query: 204 tags.append(game.c.eco >= tag_query["eco_from"]) 205 206 if "eco_to" in tag_query: 207 tags.append(game.c.eco <= tag_query["eco_to"]) 208 209 if "annotator" in tag_query: 210 tags.append(annotator.c.name.like("%%%s%%" % tag_query["annotator"])), 211 212 if "variant" in tag_query: 213 tags.append(game.c.variant == int(tag_query["variant"])), 214 215 if "result" in tag_query: 216 tags.append(game.c.result == reprResult.index(tag_query["result"])), 217 218 if "date_from" in tag_query: 219 tags.append(game.c.date >= tag_query["date_from"]) 220 221 if "date_to" in tag_query: 222 # When date_to is not given as full date we have to prepare 223 # date_to filled with some "?" to get correct query results 224 # because for example "2018.??.??" is greater than "2018" 225 date_to = tag_query["date_to"] 226 y, m, d = parseDateTag(date_to) 227 y = "%04d" % y if y is not None else "????" 228 m = "%02d" % m if m is not None else "??" 229 d = "%02d" % d if d is not None else "??" 230 date_to = "%s.%s.%s" % (y, m, d) 231 232 tags.append(game.c.date <= date_to) 233 234 if "elo_from" in tag_query: 235 tags.append(game.c.white_elo >= tag_query["elo_from"]) 236 tags.append(game.c.black_elo >= tag_query["elo_from"]) 237 238 # white_elo and black_elo are String(4) in game table 239 # we need to prevent selecting games where a player elo is 999 or 400 240 # when tag_query["elo_from"] is for example was set to 2200 241 tags.append(game.c.white_elo < "4") 242 tags.append(game.c.black_elo < "4") 243 244 if "elo_to" in tag_query: 245 tags.append(game.c.white_elo <= tag_query["elo_to"]) 246 tags.append(game.c.black_elo <= tag_query["elo_to"]) 247 248 self.where_tags = and_(*tags) 249 else: 250 self.where_tags = None 251 252 def build_where_offs8(self, offset_list): 253 if offset_list is not None and len(offset_list) > 0: 254 self.where_offs8 = game.c.offset8.in_(offset_list) 255 else: 256 self.where_offs8 = None 257 258 def build_where_offs(self, offset_list): 259 if offset_list is not None and len(offset_list) > 0: 260 self.where_offs = game.c.offset.in_(offset_list) 261 else: 262 self.where_offs = None 263 264 def build_query(self): 265 self.query = self.select 266 267 if self.where_tags is not None: 268 self.query = self.query.where(self.where_tags) 269 270 if self.where_offs8 is not None: 271 self.query = self.query.where(self.where_offs8) 272 273 if self.where_offs is not None: 274 self.query = self.query.where(self.where_offs) 275 276 def get_records(self, last_seen, limit): 277 self.build_query() 278 # we use .where() to implement pagination because .offset() doesn't scale on big tables 279 # http://sqlite.org/cvstrac/wiki?p=ScrollingCursor 280 # https://stackoverflow.com/questions/21082956/sqlite-scrolling-cursor-how-to-scroll-correctly-with-duplicate-names 281 if self.is_desc: 282 query = self.query.where(or_(self.order_cols[0] < last_seen[0], 283 and_(self.order_cols[0] == last_seen[0], 284 self.order_cols[1] < last_seen[1])) 285 ).order_by(self.order_cols[0].desc(), self.order_cols[1].desc()).limit(limit) 286 else: 287 query = self.query.where(or_(self.order_cols[0] > last_seen[0], 288 and_(self.order_cols[0] == last_seen[0], 289 self.order_cols[1] > last_seen[1])) 290 ).order_by(*self.order_cols).limit(limit) 291 292 # log.debug(self.engine.execute(Explain(query)).fetchall(), extra={"task": "SQL"}) 293 294 result = self.engine.execute(query) 295 records = result.fetchall() 296 297 return records 298 299 def get_offsets_for_tags(self, last_seen): 300 query = self.select_offsets.where(self.where_tags).where(game.c.offset > last_seen[1]) 301 result = self.engine.execute(query) 302 return [rec[0] for rec in result.fetchall()] 303 304 def get_info(self, rec): 305 where = and_(game.c.source_id == source.c.id, game.c.id == rec["Id"]) 306 result = self.engine.execute(select([source.c.info]).where(where)).first() 307 308 if result is None: 309 return None 310 else: 311 return result[0] 312 313 def get_exta_tags(self, rec): 314 return self.engine.execute(select([tag_game]).where(tag_game.c.game_id == rec["Id"])) 315