add dataset updates with doc updates

server/storage/EdgeDB.py

@@ -1,7 +1,13 @@
 from server.GameBoard import GameBoard
+from server.Dataset import Dataset
 
 from datetime import datetime
-import gel, json, time
+import json, time
+
+try:
+  import gel as _gel  # type: ignore[import-not-found]
+except ImportError:  # pragma: no cover
+  _gel = None
 
 class EdgeDB:
   def __init__(self, database:str=None, tls_security:str='insecure', **kwargs):
@@ -10,16 +16,20 @@ class EdgeDB:
     self._connect()
 
   def _connect(self):
-    self.client = gel.create_client(
-      tls_security=self.tls_security,
-      database=self.database
+    if _gel is None:
+      raise ImportError("The 'gel' package is required to use EdgeDB storage")
+    self.client = _gel.create_client(
+      tls_security=self.tls_security, database=self.database
     )
 
   def run_query_with_reconnection(self, function, *args, **kwargs):
     while True:
       try:
         return function(*args, **kwargs)
-      except gel.errors.ClientConnectionFailedError:
+      except Exception as error:
+        if error.__class__.__name__ != "ClientConnectionFailedError":
+          raise
+
         self._connect()
         time.sleep(0.5)
 
@@ -39,9 +49,21 @@ class EdgeDB:
     data = []
     moves = game_board.turns
     snake_calulations = [[calc for calc in ele["data"]] for ele in game_board.snake_class.get_history() ]
+    labels_by_turn = Dataset(game_board).labels_by_turn()
 
     for i in range(len(moves)):
-      data.append({"turn": moves[i]["turn"], "move": moves[i]["move"], "game_board": moves[i]["game_board"], "calculations": snake_calulations[i]})
+      calculations = snake_calulations[i] if i < len(snake_calulations) else []
+      calculations.append({
+        "dataset": {
+            "is_good_move": labels_by_turn.get(moves[i]["turn"], False)
+          }
+      })
+      data.append({
+        "turn": moves[i]["turn"],
+        "move": moves[i]["move"],
+        "game_board": moves[i]["game_board"],
+        "calculations": calculations,
+      })
 
     return data
 
@@ -91,9 +113,20 @@ class EdgeDB:
       created_at=datetime.fromtimestamp(game_board.now_date.timestamp(), game_board.now_date.astimezone().tzinfo),
       turns=game_board.turn,
       map=game_board.map if game_board.map else "standard",
-      winner=', '.join(game_board.winner_snake_names) if game_board.winner_snake_names else "",
-      moves=[ tuple([x["turn"], x["move"], json.dumps(x["game_board"]), [ json.dumps(ele) for ele in x["calculations"] ] ]) for x in self.create_moves_with_calculations(game_board) ],
-
+      winner=', '.join(game_board.winner_snake_names)
+      if game_board.winner_snake_names
+      else "",
+      moves=[
+        tuple(
+          [
+            x["turn"],
+            x["move"],
+            json.dumps(x["game_board"]),
+            [json.dumps(ele) for ele in x["calculations"]],
+          ]
+        )
+        for x in self.create_moves_with_calculations(game_board)
+      ],
       game_type=game_type["name"],
       is_ladder=game_type["is_ladder"],