from core.models.Model import Model from dataclasses import dataclass, field from dataclasses_json import config, Exclude from typing import Optional _table_name: str = 'locations' _table_definition: str = """ 'id' int UNIQUE, 'country_code' varchar, 'country_name' varchar, 'code' varchar, 'name' varchar, 'time_zone' varchar, UNIQUE(code, country_code) """ @dataclass class Location(Model): country_code: str code: str id: Optional[int] = field( default=None, metadata=config(exclude=Exclude.ALWAYS) ) country_name: Optional[str] = field( default=None, metadata=config(exclude=Exclude.ALWAYS) ) name: Optional[str] = field( default=None, metadata=config(exclude=Exclude.ALWAYS) ) time_zone: Optional[str] = None available: Optional[bool] = field( default=False, metadata=config(exclude=Exclude.ALWAYS) ) def __post_init__(self): self.available = self.exists(self.country_code, self.code) def is_available(self): return self.exists(self.country_code, self.code) @staticmethod def find_by_id(id: int): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) return Model._query_one('SELECT * FROM locations WHERE id = ? LIMIT 1', Location.factory, [id]) @staticmethod def find(country_code: str, code: str): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) return Model._query_one('SELECT * FROM locations WHERE country_code = ? AND code = ? LIMIT 1', Location.factory, [country_code, code]) @staticmethod def exists(country_code: str, code: str): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) return Model._query_exists('SELECT * FROM locations WHERE country_code = ? AND code = ?', [country_code, code]) @staticmethod def all(): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) return Model._query_all('SELECT * FROM locations', Location.factory) @staticmethod def truncate(): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition, drop_existing=True) @staticmethod def save_many(locations): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) Model._insert_many('INSERT INTO locations VALUES(?, ?, ?, ?, ?, ?)', Location.tuple_factory, locations) @staticmethod def factory(cursor, row): local_fields = [column[0] for column in cursor.description] return Location(**{key: value for key, value in zip(local_fields, row)}) @staticmethod def tuple_factory(location): return location.id, location.country_code, location.country_name, location.code, location.name, location.time_zone