From 545ceac0b8db577e4ab2dbe5a28bd89c8169f1ca Mon Sep 17 00:00:00 2001
From: Vladless <curov2014@yandex.ru>
Date: Wed, 13 Nov 2024 18:51:25 +0300
Subject: [PATCH] custom trial time

---
 database.py                |  2 +-
 handlers/keys/trial_key.py |  4 ++--
 models.py                  | 44 ++++++++++++++++++++++++++------------
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/database.py b/database.py
index 01bc27e3..1459cecc 100644
--- a/database.py
+++ b/database.py
@@ -3,7 +3,7 @@ from datetime import datetime
 import asyncpg
 from loguru import logger
 
-from config import DATABASE_URL, BONUS_PERCENT
+from config import BONUS_PERCENT, DATABASE_URL
 
 
 async def init_db():
diff --git a/handlers/keys/trial_key.py b/handlers/keys/trial_key.py
index 048d2946..c4138b93 100644
--- a/handlers/keys/trial_key.py
+++ b/handlers/keys/trial_key.py
@@ -6,7 +6,7 @@ import asyncpg
 from py3xui import AsyncApi
 
 from client import add_client
-from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL, PUBLIC_LINK, SERVERS
+from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL, PUBLIC_LINK, SERVERS, TRIAL_TIME
 from database import store_key
 from handlers.texts import INSTRUCTIONS
 from handlers.utils import generate_random_email
@@ -39,7 +39,7 @@ async def generate_and_store_keys(
     conn = await asyncpg.connect(DATABASE_URL)
     try:
         current_time = datetime.utcnow()
-        expiry_time = current_time + timedelta(days=1, hours=3)
+        expiry_time = current_time + timedelta(days={TRIAL_TIME}, hours=3)
         expiry_timestamp = int(expiry_time.timestamp() * 1000)
 
         tasks = []
diff --git a/models.py b/models.py
index f3ea708d..ee997b3a 100644
--- a/models.py
+++ b/models.py
@@ -1,54 +1,70 @@
-from sqlalchemy import Column, BigInteger, Float, Text, Boolean, ForeignKey
+from sqlalchemy import BigInteger, Boolean, Column, Float, ForeignKey, Text
 from sqlalchemy.orm import declarative_base, relationship
 
 Base = declarative_base()
 
+
 class Connection(Base):
     """
     Модель для таблицы connections
     """
-    __tablename__ = 'connections'
+
+    __tablename__ = "connections"
 
     tg_id = Column(BigInteger, primary_key=True, nullable=False)
     balance = Column(Float, nullable=False, default=0.0)
     trial = Column(BigInteger, nullable=False, default=0)
 
     # Связь с ключами и рефералами
-    keys = relationship('Key', back_populates='connection')
-    referrals_received = relationship('Referral', foreign_keys='Referral.referred_tg_id', back_populates='referred')
-    referrals_sent = relationship('Referral', foreign_keys='Referral.referrer_tg_id', back_populates='referrer')
+    keys = relationship("Key", back_populates="connection")
+    referrals_received = relationship(
+        "Referral", foreign_keys="Referral.referred_tg_id", back_populates="referred"
+    )
+    referrals_sent = relationship(
+        "Referral", foreign_keys="Referral.referrer_tg_id", back_populates="referrer"
+    )
 
 
 class Key(Base):
     """
     Модель для таблицы keys
     """
-    __tablename__ = 'keys'
 
-    tg_id = Column(BigInteger, ForeignKey('connections.tg_id'), primary_key=True, nullable=False)
+    __tablename__ = "keys"
+
+    tg_id = Column(
+        BigInteger, ForeignKey("connections.tg_id"), primary_key=True, nullable=False
+    )
     client_id = Column(Text, primary_key=True, nullable=False)
     email = Column(Text, nullable=False)
     created_at = Column(BigInteger, nullable=False)
     expiry_time = Column(BigInteger, nullable=False)
     key = Column(Text, nullable=False)
-    server_id = Column(Text, nullable=False, default='server1')
+    server_id = Column(Text, nullable=False, default="server1")
     notified = Column(Boolean, nullable=False, default=False)
     notified_24h = Column(Boolean, nullable=False, default=False)
 
     # Связь с подключением
-    connection = relationship('Connection', back_populates='keys')
+    connection = relationship("Connection", back_populates="keys")
 
 
 class Referral(Base):
     """
     Модель для таблицы referrals
     """
-    __tablename__ = 'referrals'
 
-    referred_tg_id = Column(BigInteger, ForeignKey('connections.tg_id'), primary_key=True, nullable=False)
-    referrer_tg_id = Column(BigInteger, ForeignKey('connections.tg_id'), nullable=False)
+    __tablename__ = "referrals"
+
+    referred_tg_id = Column(
+        BigInteger, ForeignKey("connections.tg_id"), primary_key=True, nullable=False
+    )
+    referrer_tg_id = Column(BigInteger, ForeignKey("connections.tg_id"), nullable=False)
     reward_issued = Column(Boolean, default=False)
 
     # Связи с подключениями
-    referred = relationship('Connection', foreign_keys=[referred_tg_id], back_populates='referrals_received')
-    referrer = relationship('Connection', foreign_keys=[referrer_tg_id], back_populates='referrals_sent') 
\ No newline at end of file
+    referred = relationship(
+        "Connection", foreign_keys=[referred_tg_id], back_populates="referrals_received"
+    )
+    referrer = relationship(
+        "Connection", foreign_keys=[referrer_tg_id], back_populates="referrals_sent"
+    )