mbot/db.py

50 lines
1.4 KiB
Python
Raw Permalink Normal View History

2024-06-08 18:16:53 +00:00
from sqlalchemy import create_engine, Column, String, BigInteger, ForeignKey, TIMESTAMP, UniqueConstraint
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship, sessionmaker
from datetime import datetime
Base = declarative_base()
class Token(Base):
__tablename__ = 'tokens'
id = Column(String(42), primary_key=True)
name = Column(String(255), nullable=True)
symbol = Column(String(32), nullable=True)
total_supply = Column(String, nullable=True)
class Pair(Base):
__tablename__ = 'pairs'
id = Column(String(42), primary_key=True)
token0 = Column(String(42), ForeignKey('tokens.id'), nullable=False)
token1 = Column(String(42), ForeignKey('tokens.id'), nullable=False)
created_at = Column(TIMESTAMP, default=datetime.utcnow)
# Ensuring the combination of token0, token1, and pool_address is unique
__table_args__ = (UniqueConstraint('token0', 'token1', name='_token_pool_uc'),)
# Relationships
token0_rel = relationship("Token", foreign_keys=[token0])
token1_rel = relationship("Token", foreign_keys=[token1])
# Engine creation
engine = create_engine('sqlite:///uniswap.db', echo=True)
# Create tables
Base.metadata.create_all(engine)
# Create a session factory
SessionLocal = sessionmaker(bind=engine)
# Dependency to get a database session
def get_session():
session = SessionLocal()
try:
yield session
finally:
session.close()