walrasquant-lib 0.4.20__py3-none-any.whl

walrasquant/backends/db_postgresql.py

@@ -0,0 +1,321 @@
+import asyncpg
+import psycopg2
+import psycopg2.extensions
+from decimal import Decimal
+from typing import Dict, Set, List, Optional, Any
+
+from walrasquant.backends.db import StorageBackend
+from walrasquant.schema import Order, Position, Balance, AccountBalance
+from walrasquant.constants import AccountType, ExchangeType, get_postgresql_config
+
+
+class PostgreSQLBackend(StorageBackend):
+    def __init__(
+        self, strategy_id: str, user_id: str, table_prefix: str, log, **kwargs
+    ):
+        super().__init__(strategy_id, user_id, table_prefix, log, **kwargs)
+        self._pg_async: Optional[asyncpg.Pool] = None
+        self._pg: Optional[psycopg2.extensions.connection] = None
+
+    @property
+    def _pool(self) -> asyncpg.Pool:
+        assert self._pg_async is not None
+        return self._pg_async
+
+    @property
+    def _sync_conn(self) -> psycopg2.extensions.connection:
+        assert self._pg is not None
+        return self._pg
+
+    async def _init_conn(self) -> None:
+        pg_config = get_postgresql_config()
+        self._pg_async = await asyncpg.create_pool(**pg_config)
+        self._pg = psycopg2.connect(**pg_config)
+
+    async def _init_table(self) -> None:
+        async with self._pool.acquire() as conn:
+            await conn.execute(f"""
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_orders (
+                    timestamp BIGINT,
+                    oid TEXT PRIMARY KEY,
+                    eid TEXT,
+                    symbol TEXT,
+                    side TEXT, 
+                    type TEXT,
+                    amount TEXT,
+                    price DOUBLE PRECISION,
+                    status TEXT,
+                    fee TEXT,
+                    fee_currency TEXT,
+                    data BYTEA
+                );
+                
+                CREATE INDEX IF NOT EXISTS idx_orders_symbol
+                ON {self.table_prefix}_orders(symbol);
+
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_positions (
+                    symbol TEXT PRIMARY KEY,
+                    exchange TEXT,
+                    side TEXT,
+                    amount TEXT,
+                    data BYTEA
+                );
+                
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_open_orders (
+                    oid TEXT PRIMARY KEY,
+                    eid TEXT,
+                    exchange TEXT,
+                    symbol TEXT
+                );
+                
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_balances (
+                    asset TEXT,
+                    account_type TEXT,
+                    free TEXT,
+                    locked TEXT,
+                    PRIMARY KEY (asset, account_type)
+                );
+                
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_params (
+                    key TEXT PRIMARY KEY,
+                    value BYTEA,
+                    timestamp BIGINT DEFAULT (EXTRACT(EPOCH FROM NOW()) * 1000)::BIGINT
+                );
+            """)
+
+    async def close(self) -> None:
+        if self._pg_async:
+            await self._pg_async.close()
+        if self._pg:
+            self._pg.close()
+
+    async def sync_orders(self, mem_orders: Dict[str, Order]) -> None:
+        async with self._pool.acquire() as conn:
+            for order in mem_orders.copy().values():
+                await conn.execute(
+                    f"INSERT INTO {self.table_prefix}_orders "
+                    "(timestamp, oid, eid, symbol, side, type, amount, price, status, fee, fee_currency, data) "
+                    "VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) "
+                    "ON CONFLICT (oid) DO UPDATE SET "
+                    "timestamp = EXCLUDED.timestamp, "
+                    "eid = EXCLUDED.eid, "
+                    "symbol = EXCLUDED.symbol, "
+                    "side = EXCLUDED.side, "
+                    "type = EXCLUDED.type, "
+                    "amount = EXCLUDED.amount, "
+                    "price = EXCLUDED.price, "
+                    "status = EXCLUDED.status, "
+                    "fee = EXCLUDED.fee, "
+                    "fee_currency = EXCLUDED.fee_currency, "
+                    "data = EXCLUDED.data",
+                    order.timestamp or 0,
+                    str(order.oid),
+                    str(order.eid),
+                    str(order.symbol),
+                    str(order.side.value) if order.side is not None else None,
+                    str(order.type.value) if order.type is not None else None,
+                    str(order.amount),
+                    float(_p)
+                    if (_p := order.price or order.average) is not None
+                    else None,
+                    str(order.status.value),
+                    str(order.fee) if order.fee is not None else None,
+                    order.fee_currency,
+                    self._encode(order),
+                )
+
+    async def sync_positions(self, mem_positions: Dict[str, Position]) -> None:
+        async with self._pool.acquire() as conn:
+            db_positions = await conn.fetch(
+                f"SELECT symbol FROM {self.table_prefix}_positions"
+            )
+            db_position_symbols = {row["symbol"] for row in db_positions}
+
+            positions_to_delete = db_position_symbols - set(mem_positions.keys())
+            if positions_to_delete:
+                for symbol in positions_to_delete:
+                    await conn.execute(
+                        f"DELETE FROM {self.table_prefix}_positions WHERE symbol = $1",
+                        str(symbol),
+                    )
+                self._log.debug(
+                    f"Deleted {len(positions_to_delete)} stale positions from database"
+                )
+
+            for symbol, position in mem_positions.copy().items():
+                await conn.execute(
+                    f"INSERT INTO {self.table_prefix}_positions "
+                    "(symbol, exchange, side, amount, data) VALUES ($1, $2, $3, $4, $5) "
+                    "ON CONFLICT (symbol) DO UPDATE SET "
+                    "exchange = EXCLUDED.exchange, "
+                    "side = EXCLUDED.side, "
+                    "amount = EXCLUDED.amount, "
+                    "data = EXCLUDED.data",
+                    str(symbol),
+                    str(position.exchange.value),
+                    str(position.side.value) if position.side else "FLAT",
+                    str(position.amount),
+                    self._encode(position),
+                )
+
+    async def sync_open_orders(
+        self,
+        mem_open_orders: Dict[ExchangeType, Set[str]],
+        mem_orders: Dict[str, Order],
+    ) -> None:
+        async with self._pool.acquire() as conn:
+            await conn.execute(f"DELETE FROM {self.table_prefix}_open_orders")
+
+            for exchange, oids in mem_open_orders.copy().items():
+                for oid in oids.copy():
+                    order = mem_orders.get(oid)
+                    if order:
+                        await conn.execute(
+                            f"INSERT INTO {self.table_prefix}_open_orders "
+                            "(oid, eid, exchange, symbol) VALUES ($1, $2, $3, $4)",
+                            str(oid),
+                            str(order.eid),
+                            str(exchange.value),
+                            str(order.symbol),
+                        )
+
+    async def sync_balances(
+        self, mem_account_balance: Dict[AccountType, AccountBalance]
+    ) -> None:
+        async with self._pool.acquire() as conn:
+            for account_type, balance in mem_account_balance.copy().items():
+                for asset, amount in balance.balances.items():
+                    await conn.execute(
+                        f"INSERT INTO {self.table_prefix}_balances "
+                        "(asset, account_type, free, locked) VALUES ($1, $2, $3, $4) "
+                        "ON CONFLICT (asset, account_type) DO UPDATE SET "
+                        "free = EXCLUDED.free, "
+                        "locked = EXCLUDED.locked",
+                        str(asset),
+                        str(account_type.value),
+                        str(amount.free),
+                        str(amount.locked),
+                    )
+
+    def get_order(
+        self,
+        oid: str,
+        mem_orders: Dict[str, Order],
+    ) -> Optional[Order]:
+        if order := mem_orders.get(oid):
+            return order
+
+        try:
+            cursor = self._sync_conn.cursor()
+            with cursor:
+                cursor.execute(
+                    f"SELECT data FROM {self.table_prefix}_orders WHERE oid = %s",
+                    (oid,),
+                )
+                row = cursor.fetchone()
+                if row:
+                    order = self._decode(row[0], Order)
+                    mem_orders[oid] = order
+                    return order
+                return None
+        except psycopg2.DatabaseError as error:
+            self._log.error(f"Error getting order from PostgreSQL: {error}")
+            return None
+
+    def get_symbol_orders(self, symbol: str) -> Set[str]:
+        cursor = self._sync_conn.cursor()
+        with cursor:
+            cursor.execute(
+                f"SELECT oid FROM {self.table_prefix}_orders WHERE symbol = %s",
+                (symbol,),
+            )
+            result = {row[0] for row in cursor.fetchall()}
+        return result
+
+    def get_all_positions(self, exchange_id: ExchangeType) -> Dict[str, Position]:
+        positions = {}
+        cursor = self._sync_conn.cursor()
+        cursor.execute(
+            f"SELECT symbol, data FROM {self.table_prefix}_positions WHERE exchange = %s",
+            (exchange_id.value,),
+        )
+        for row in cursor.fetchall():
+            position = self._decode(row[1], Position)
+            if position.side:
+                positions[position.symbol] = position
+        cursor.close()
+        return positions
+
+    def get_all_balances(self, account_type: AccountType) -> List[Balance]:
+        balances = []
+        cursor = self._sync_conn.cursor()
+        cursor.execute(
+            f"SELECT asset, free, locked FROM {self.table_prefix}_balances WHERE account_type = %s",
+            (account_type.value,),
+        )
+        for row in cursor.fetchall():
+            balances.append(
+                Balance(asset=row[0], free=Decimal(row[1]), locked=Decimal(row[2]))
+            )
+        cursor.close()
+        return balances
+
+    async def sync_params(self, mem_params: Dict[str, Any]) -> None:
+        import msgspec
+
+        async with self._pool.acquire() as conn:
+            for key, value in mem_params.copy().items():
+                try:
+                    serialized_value = self._encode_param(value)
+                    await conn.execute(
+                        f"INSERT INTO {self.table_prefix}_params "
+                        "(key, value) VALUES ($1, $2) "
+                        "ON CONFLICT (key) DO UPDATE SET "
+                        "value = EXCLUDED.value, "
+                        "timestamp = DEFAULT",
+                        str(key),
+                        serialized_value,
+                    )
+                except msgspec.EncodeError as e:
+                    self._log.error(f"Error serializing parameter {key}: {e}")
+                except Exception as e:
+                    self._log.error(f"Error storing parameter {key}: {e}")
+
+    def get_param(self, key: str, default: Any = None) -> Any:
+        import msgspec
+
+        try:
+            cursor = self._sync_conn.cursor()
+            cursor.execute(
+                f"SELECT value FROM {self.table_prefix}_params WHERE key = %s", (key,)
+            )
+            row = cursor.fetchone()
+            cursor.close()
+
+            if row:
+                try:
+                    return self._decode_param(row[0])
+                except msgspec.DecodeError as e:
+                    self._log.error(f"Error deserializing parameter {key}: {e}")
+                    return default
+            return default
+        except Exception as e:
+            self._log.error(f"Error getting parameter {key}: {e}")
+            return default
+
+    def get_all_params(self) -> Dict[str, Any]:
+        import msgspec
+
+        params = {}
+        try:
+            cursor = self._sync_conn.cursor()
+            cursor.execute(f"SELECT key, value FROM {self.table_prefix}_params")
+            for row in cursor.fetchall():
+                try:
+                    params[row[0]] = self._decode_param(row[1])
+                except msgspec.DecodeError as e:
+                    self._log.error(f"Error deserializing parameter {row[0]}: {e}")
+            cursor.close()
+        except Exception as e:
+            self._log.error(f"Error getting all parameters: {e}")
+        return params

walrasquant/backends/db_sqlite.py

@@ -0,0 +1,310 @@
+import aiosqlite
+import sqlite3
+from pathlib import Path
+from typing import Dict, Set, List, Optional, Any
+
+from walrasquant.backends.db import StorageBackend
+from walrasquant.schema import Order, Position, Balance, AccountBalance
+from walrasquant.constants import AccountType, ExchangeType
+
+
+class SQLiteBackend(StorageBackend):
+    def __init__(
+        self,
+        strategy_id: str,
+        user_id: str,
+        table_prefix: str,
+        log,
+        db_path: str = ".keys/cache.db",
+        **kwargs,
+    ):
+        super().__init__(strategy_id, user_id, table_prefix, log, **kwargs)
+        self.db_path = db_path
+        self._db_async: Optional[aiosqlite.Connection] = None
+        self._db: Optional[sqlite3.Connection] = None
+
+    @property
+    def _async_conn(self) -> aiosqlite.Connection:
+        assert self._db_async is not None
+        return self._db_async
+
+    @property
+    def _sync_db(self) -> sqlite3.Connection:
+        assert self._db is not None
+        return self._db
+
+    async def _init_conn(self) -> None:
+        db_path = Path(self.db_path)
+        db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._db_async = await aiosqlite.connect(str(db_path))
+        self._db = sqlite3.connect(str(db_path))
+
+    async def _init_table(self) -> None:
+        async with self._async_conn.cursor() as cursor:
+            await cursor.executescript(f"""
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_orders (
+                    timestamp INTEGER,
+                    oid TEXT PRIMARY KEY,
+                    eid TEXT,
+                    symbol TEXT,
+                    side TEXT, 
+                    type TEXT,
+                    amount TEXT,
+                    price REAL,
+                    status TEXT,
+                    fee TEXT,
+                    fee_currency TEXT,
+                    data BLOB
+                );
+                
+                CREATE INDEX IF NOT EXISTS idx_orders_symbol
+                ON {self.table_prefix}_orders(symbol);
+
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_positions (
+                    symbol PRIMARY KEY,
+                    exchange TEXT,
+                    side TEXT,
+                    amount TEXT,
+                    data BLOB
+                );
+                
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_open_orders (
+                    oid PRIMARY KEY,
+                    eid TEXT,
+                    exchange TEXT,
+                    symbol TEXT
+                );
+                
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_balances (
+                    asset TEXT,
+                    account_type TEXT,
+                    free TEXT,
+                    locked TEXT,
+                    PRIMARY KEY (asset, account_type)
+                );
+                
+                CREATE TABLE IF NOT EXISTS {self.table_prefix}_params (
+                    key TEXT PRIMARY KEY,
+                    value BLOB,
+                    timestamp INTEGER DEFAULT (strftime('%s', 'now') * 1000)
+                );
+            """)
+            await self._async_conn.commit()
+
+    def _get_cursor(self):
+        """Get a thread-safe SQLite cursor, creating new connection if needed."""
+        try:
+            return self._sync_db.cursor()
+        except sqlite3.ProgrammingError:
+            # SQLite connection is from different thread, create new one
+            db_path = Path(self.db_path)
+            return sqlite3.connect(str(db_path)).cursor()
+
+    async def close(self) -> None:
+        if self._db_async:
+            await self._db_async.close()
+        if self._db:
+            self._db.close()
+
+    async def sync_orders(self, mem_orders: Dict[str, Order]) -> None:
+        async with self._async_conn.cursor() as cursor:
+            for order in mem_orders.copy().values():
+                await cursor.execute(
+                    f"INSERT OR REPLACE INTO {self.table_prefix}_orders "
+                    "(timestamp, oid, eid, symbol, side, type, amount, price, status, fee, fee_currency, data) "
+                    "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+                    (
+                        order.timestamp,
+                        order.oid,
+                        order.eid,
+                        order.symbol,
+                        order.side.value if order.side is not None else None,
+                        order.type.value if order.type is not None else None,
+                        str(order.amount),
+                        order.price or order.average,
+                        order.status.value,
+                        str(order.fee) if order.fee is not None else None,
+                        order.fee_currency,
+                        self._encode(order),
+                    ),
+                )
+            await self._async_conn.commit()
+
+    async def sync_positions(self, mem_positions: Dict[str, Position]) -> None:
+        async with self._async_conn.cursor() as cursor:
+            await cursor.execute(f"SELECT symbol FROM {self.table_prefix}_positions")
+            db_positions = {row[0] for row in await cursor.fetchall()}
+
+            positions_to_delete = db_positions - set(mem_positions.keys())
+            if positions_to_delete:
+                await cursor.executemany(
+                    f"DELETE FROM {self.table_prefix}_positions WHERE symbol = ?",
+                    [(symbol,) for symbol in positions_to_delete],
+                )
+                self._log.debug(
+                    f"Deleted {len(positions_to_delete)} stale positions from database"
+                )
+
+            for symbol, position in mem_positions.copy().items():
+                await cursor.execute(
+                    f"INSERT OR REPLACE INTO {self.table_prefix}_positions "
+                    "(symbol, exchange, side, amount, data) VALUES (?, ?, ?, ?, ?)",
+                    (
+                        symbol,
+                        position.exchange.value,
+                        position.side.value if position.side else "FLAT",
+                        str(position.amount),
+                        self._encode(position),
+                    ),
+                )
+            await self._async_conn.commit()
+
+    async def sync_open_orders(
+        self,
+        mem_open_orders: Dict[ExchangeType, Set[str]],
+        mem_orders: Dict[str, Order],
+    ) -> None:
+        async with self._async_conn.cursor() as cursor:
+            await cursor.execute(f"DELETE FROM {self.table_prefix}_open_orders")
+
+            for exchange, oids in mem_open_orders.copy().items():
+                for oid in oids.copy():
+                    order = mem_orders.get(oid)
+                    if order:
+                        await cursor.execute(
+                            f"INSERT INTO {self.table_prefix}_open_orders "
+                            "(oid, eid, exchange, symbol) VALUES (?, ?, ?, ?)",
+                            (oid, order.eid, exchange.value, order.symbol),
+                        )
+            await self._async_conn.commit()
+
+    async def sync_balances(
+        self, mem_account_balance: Dict[AccountType, AccountBalance]
+    ) -> None:
+        async with self._async_conn.cursor() as cursor:
+            for account_type, balance in mem_account_balance.copy().items():
+                for asset, amount in balance.balances.items():
+                    await cursor.execute(
+                        f"INSERT OR REPLACE INTO {self.table_prefix}_balances "
+                        "(asset, account_type, free, locked) VALUES (?, ?, ?, ?)",
+                        (
+                            asset,
+                            account_type.value,
+                            str(amount.free),
+                            str(amount.locked),
+                        ),
+                    )
+            await self._async_conn.commit()
+
+    def get_order(
+        self,
+        oid: str,
+        mem_orders: Dict[str, Order],
+    ) -> Optional[Order]:
+        if order := mem_orders.get(oid):
+            return order
+
+        cursor = self._get_cursor()
+        try:
+            cursor.execute(
+                f"SELECT data FROM {self.table_prefix}_orders WHERE oid = ?",
+                (oid,),
+            )
+            if row := cursor.fetchone():
+                order = self._decode(row[0], Order)
+                mem_orders[oid] = order
+                return order
+            return None
+
+        except sqlite3.Error as e:
+            self._log.error(f"Error getting order from SQLite: {e}")
+            return None
+
+    def get_symbol_orders(self, symbol: str) -> Set[str]:
+        cursor = self._get_cursor()
+        cursor.execute(
+            f"SELECT oid FROM {self.table_prefix}_orders WHERE symbol = ?",
+            (symbol,),
+        )
+        return {row[0] for row in cursor.fetchall()}
+
+    def get_all_positions(self, exchange_id: ExchangeType) -> Dict[str, Position]:
+        positions = {}
+        cursor = self._get_cursor()
+        cursor.execute(
+            f"SELECT symbol, data FROM {self.table_prefix}_positions WHERE exchange = ?",
+            (exchange_id.value,),
+        )
+        for row in cursor.fetchall():
+            position = self._decode(row[1], Position)
+            if position.side:
+                positions[position.symbol] = position
+        return positions
+
+    def get_all_balances(self, account_type: AccountType) -> List[Balance]:
+        from decimal import Decimal
+
+        balances = []
+        cursor = self._get_cursor()
+        cursor.execute(
+            f"SELECT asset, free, locked FROM {self.table_prefix}_balances WHERE account_type = ?",
+            (account_type.value,),
+        )
+        for row in cursor.fetchall():
+            balances.append(
+                Balance(asset=row[0], free=Decimal(row[1]), locked=Decimal(row[2]))
+            )
+        return balances
+
+    async def sync_params(self, mem_params: Dict[str, Any]) -> None:
+        import msgspec
+
+        async with self._async_conn.cursor() as cursor:
+            for key, value in mem_params.copy().items():
+                try:
+                    serialized_value = self._encode_param(value)
+                    await cursor.execute(
+                        f"INSERT OR REPLACE INTO {self.table_prefix}_params "
+                        "(key, value) VALUES (?, ?)",
+                        (key, serialized_value),
+                    )
+                except msgspec.EncodeError as e:
+                    self._log.error(f"Error serializing parameter {key}: {e}")
+                except sqlite3.Error as e:
+                    self._log.error(f"Error storing parameter {key}: {e}")
+            await self._async_conn.commit()
+
+    def get_param(self, key: str, default: Any = None) -> Any:
+        import msgspec
+
+        try:
+            cursor = self._get_cursor()
+            cursor.execute(
+                f"SELECT value FROM {self.table_prefix}_params WHERE key = ?", (key,)
+            )
+            if row := cursor.fetchone():
+                try:
+                    return self._decode_param(row[0])
+                except msgspec.DecodeError as e:
+                    self._log.error(f"Error deserializing parameter {key}: {e}")
+                    return default
+            return default
+        except sqlite3.Error as e:
+            self._log.error(f"Error getting parameter {key}: {e}")
+            return default
+
+    def get_all_params(self) -> Dict[str, Any]:
+        import msgspec
+
+        params = {}
+        try:
+            cursor = self._get_cursor()
+            cursor.execute(f"SELECT key, value FROM {self.table_prefix}_params")
+            for row in cursor.fetchall():
+                try:
+                    params[row[0]] = self._decode_param(row[1])
+                except msgspec.DecodeError as e:
+                    self._log.error(f"Error deserializing parameter {row[0]}: {e}")
+        except sqlite3.Error as e:
+            self._log.error(f"Error getting all parameters: {e}")
+        return params

walrasquant/base/__init__.py

@@ -0,0 +1,24 @@
+from walrasquant.base.exchange import ExchangeManager
+from walrasquant.base.ws_client import WSClient
+from walrasquant.base.api_client import ApiClient
+from walrasquant.base.oms import OrderManagementSystem
+from walrasquant.base.ems import ExecutionManagementSystem
+from walrasquant.base.sms import SubscriptionManagementSystem
+from walrasquant.base.connector import (
+    PublicConnector,
+    PrivateConnector,
+)
+from walrasquant.base.retry import RetryManager
+
+
+__all__ = [
+    "ExchangeManager",
+    "WSClient",
+    "ApiClient",
+    "OrderManagementSystem",
+    "ExecutionManagementSystem",
+    "PublicConnector",
+    "SubscriptionManagementSystem",
+    "PrivateConnector",
+    "RetryManager",
+]

walrasquant/base/api_client.py

@@ -0,0 +1,46 @@
+from abc import ABC
+import aiohttp
+import nexuslog as logging
+from typing import Optional, Any, Dict
+from walrasquant.core.nautilius_core import LiveClock
+from walrasquant.constants import RateLimiter
+from walrasquant.base.retry import RetryManager
+
+
+class ApiClient(ABC):
+    def __init__(
+        self,
+        clock: LiveClock,
+        api_key: Optional[str] = None,
+        secret: Optional[str] = None,
+        timeout: int = 10,
+        rate_limiter: Optional[RateLimiter] = None,
+        retry_manager: Optional[RetryManager] = None,
+    ):
+        self._api_key = api_key
+        self._secret = secret
+        self._timeout = timeout
+        self._log = logging.getLogger(name=type(self).__name__)
+        self._session: Optional[aiohttp.ClientSession] = None
+        self._clock = clock
+        self._limiter = rate_limiter
+        self._retry_manager: Optional[RetryManager] = retry_manager
+        self._base_url: str | None = None
+
+    async def _init_session(self, base_url: str | None = None):
+        if self._session is None:
+            kwargs: Dict[str, Any] = {
+                "timeout": aiohttp.ClientTimeout(total=self._timeout),
+            }
+            if base_url:
+                kwargs["base_url"] = base_url
+            self._session = aiohttp.ClientSession(**kwargs)
+
+    def _get_rate_limit_cost(self, cost: int = 1):
+        return cost
+
+    async def close_session(self):
+        """Close the session"""
+        if self._session:
+            await self._session.close()
+            self._session = None