1from sqlalchemy.exc import SQLAlchemyError
 2from sqlalchemy.orm import Session
 3
 4from .base import Base, SessionLocal, engine  # noqa
 5
 6from .crud import (create_admin, create_notification_reminder,  # noqa
 7                   create_user, delete_notification_reminder, get_admin,
 8                   get_admins, get_jwt_secret_key, get_notification_reminder,
 9                   assure_node_inbounds, get_system_usage,
10                   get_tls_certificate, get_user, get_user_by_id, get_users,
11                   get_users_count, remove_admin, remove_user, revoke_user_sub,
12                   set_owner, update_admin, update_user, update_user_status,
13                   update_user_sub, start_user_expire)
14
15from .models import JWT, System, User  # noqa
16
17
18class GetDB:  # Context Manager
19    def __init__(self):
20        self.db = SessionLocal()
21
22    def __enter__(self):
23        return self.db
24
25    def __exit__(self, _, exc_value, traceback):
26        if isinstance(exc_value, SQLAlchemyError):
27            self.db.rollback()  # rollback on exception
28
29        self.db.close()
30
31
32__all__ = [
33    "get_user",
34    "get_user_by_id",
35    "get_users",
36    "get_users_count",
37    "create_user",
38    "remove_user",
39    "update_user",
40    "update_user_status",
41    "start_user_expire",
42    "update_user_sub",
43    "revoke_user_sub",
44    "set_owner",
45    "get_system_usage",
46    "get_jwt_secret_key",
47    "get_tls_certificate",
48    "get_admin",
49    "create_admin",
50    "update_admin",
51    "remove_admin",
52    "get_admins",
53
54    "create_notification_reminder",
55    "get_notification_reminder",
56    "delete_notification_reminder",
57
58    "GetDB",
59
60    "User",
61    "System",
62    "JWT",
63
64    "Base",
65    "Session",
66]