snake-python/server/storage/EdgeDB.py

from server.GameBoard import GameBoard
from server.Dataset import Dataset

from datetime import datetime
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):
    self.database = database
    self.tls_security = tls_security
    self._connect()

  def _connect(self):
    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 Exception as error:
        if error.__class__.__name__ != "ClientConnectionFailedError":
          raise

        self._connect()
        time.sleep(0.5)

  def insert_game_type(self, name:str, is_ladder:bool):
    return self.run_query_with_reconnection(
      self.client.query_required_single,
      """
      insert GameType {
        name := <str>$name,
        is_ladder := <bool>$is_ladder
      }""",
      name=name,
      is_ladder=is_ladder
    )

  def create_moves_with_calculations(self, game_board:GameBoard):
    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)):
      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

  async def insert(self, game_board:GameBoard):
    game_type = game_board.get_type_of_game()

    self.run_query_with_reconnection(
      self.client.query,
      """
      insert GameBoard {
        id := <uuid>$id,
        created_at := <datetime>$created_at,
        turns := <int32>$turns,
        map := <str>$map,
        winner := <str>$winner,
        moves := (
          with input_data := <array <tuple <turn: int32, `move`: str, game_board: json, calculations: array<json> >>>$moves
          for data in array_unpack(input_data)
            insert Moves {
              turn := data.turn,
              snake_move := data.`move`,
              game_board := data.game_board,
              calculations := data.calculations
            }
        ),
        type := (
          insert GameType {
            name := <str>$game_type,
            is_ladder := <bool>$is_ladder
          } unless conflict on (.name, .is_ladder) else GameType
        ),
        ruleset := (
          insert Ruleset {
            name := <str>$ruleset,
            version := <str>$version,
            settings := to_json(<str>$settings)
          } unless conflict on (.name, .version, .settings) else Ruleset
        ),
        snake := (
          insert Snake {
            type := <str>$snake_type
          } unless conflict on .type else Snake
        )
      }""",

      id=game_board.id,
      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)
      ],
      game_type=game_type["name"],
      is_ladder=game_type["is_ladder"],

      ruleset=game_board.ruleset["name"],
      version=game_board.ruleset["version"],
      settings=json.dumps(game_board.ruleset["settings"]),

      snake_type=game_board.snake_class.__class__.__name__,
    )

  async def save(self, game_board:GameBoard):
    await self.insert(game_board)

  def __del__(self):
    self.client.close()

  def cleanup(self):
    return self.run_query_with_reconnection(
      self.client.query_json,
      """
      delete Moves { };

      with gameboard := (delete GameBoard filter .turns < <std::int32>"200" or .is_winner_me = <std::bool>"false")
      select gameboard {id, url, winner, turns, type: { is_ladder, name } } order by .turns desc;

      """
    )