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