freqtrade_origin/freqtrade/persistence/models.py

85 lines
2.9 KiB
Python
Raw Normal View History

2018-03-02 15:22:00 +00:00
"""
This module contains the class to persist trades into SQLite
"""
import logging
import threading
from contextvars import ContextVar
from typing import Any, Dict, Final, Optional
2017-05-12 17:11:56 +00:00
from sqlalchemy import create_engine, inspect
2018-06-07 19:35:57 +00:00
from sqlalchemy.exc import NoSuchModuleError
2023-02-28 06:03:22 +00:00
from sqlalchemy.orm import scoped_session, sessionmaker
2017-11-09 22:45:22 +00:00
from sqlalchemy.pool import StaticPool
2017-05-12 17:11:56 +00:00
from freqtrade.exceptions import OperationalException
2023-03-01 11:16:56 +00:00
from freqtrade.persistence.base import ModelBase
2023-04-08 08:02:38 +00:00
from freqtrade.persistence.key_value_store import _KeyValueStoreModel
2020-08-13 06:33:46 +00:00
from freqtrade.persistence.migrations import check_migrate
from freqtrade.persistence.pairlock import PairLock
from freqtrade.persistence.trade_model import Order, Trade
2018-06-07 19:35:57 +00:00
2020-09-28 17:39:41 +00:00
logger = logging.getLogger(__name__)
2019-09-10 07:42:45 +00:00
REQUEST_ID_CTX_KEY: Final[str] = 'request_id'
_request_id_ctx_var: ContextVar[Optional[str]] = ContextVar(REQUEST_ID_CTX_KEY, default=None)
def get_request_or_thread_id() -> Optional[str]:
"""
Helper method to get either async context (for fastapi requests), or thread id
"""
id = _request_id_ctx_var.get()
if id is None:
# when not in request context - use thread id
id = str(threading.current_thread().ident)
return id
2018-06-23 13:27:29 +00:00
_SQL_DOCS_URL = 'http://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls'
2017-05-12 17:11:56 +00:00
2022-05-19 04:45:20 +00:00
def init_db(db_url: str) -> None:
"""
2021-11-09 18:22:29 +00:00
Initializes this module with the given config,
registers all known command handlers
and starts polling for message updates
:param db_url: Database to use
:return: None
"""
2023-02-14 10:04:09 +00:00
kwargs: Dict[str, Any] = {}
2022-02-22 18:39:55 +00:00
if db_url == 'sqlite:///':
raise OperationalException(
f'Bad db-url {db_url}. For in-memory database, please use `sqlite://`.')
if db_url == 'sqlite://':
kwargs.update({
'poolclass': StaticPool,
2021-05-23 06:56:41 +00:00
})
# Take care of thread ownership
if db_url.startswith('sqlite://'):
kwargs.update({
'connect_args': {'check_same_thread': False},
})
2018-06-07 19:35:57 +00:00
try:
2021-04-13 15:34:20 +00:00
engine = create_engine(db_url, future=True, **kwargs)
2018-06-07 19:35:57 +00:00
except NoSuchModuleError:
raise OperationalException(f"Given value for db_url: '{db_url}' "
f"is no valid database URL! (See {_SQL_DOCS_URL})")
2018-06-07 19:35:57 +00:00
2019-10-29 13:26:03 +00:00
# https://docs.sqlalchemy.org/en/13/orm/contextual.html#thread-local-scope
# Scoped sessions proxy requests to the appropriate thread-local session.
# Since we also use fastAPI, we need to make it aware of the request id, too
Trade.session = scoped_session(sessionmaker(
bind=engine, autoflush=False), scopefunc=get_request_or_thread_id)
2023-03-15 20:12:06 +00:00
Order.session = Trade.session
PairLock.session = Trade.session
2023-04-08 08:02:38 +00:00
_KeyValueStoreModel.session = Trade.session
2020-10-17 09:28:34 +00:00
previous_tables = inspect(engine).get_table_names()
2023-02-16 07:04:45 +00:00
ModelBase.metadata.create_all(engine)
check_migrate(engine, decl_base=ModelBase, previous_tables=previous_tables)