Coverage for app/backend/src/couchers/db.py: 83%
97 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 22:32 +0000
1import functools
2import inspect
3import logging
4import os
5from collections.abc import Generator, Sequence
6from contextlib import contextmanager
7from os import getpid
8from threading import get_ident
9from typing import TYPE_CHECKING
11from alembic import command
12from alembic.config import Config
13from geoalchemy2 import WKBElement
14from opentelemetry import trace
15from sqlalchemy import Engine, Row, Subquery, create_engine, select, text, true
16from sqlalchemy.dialects import registry
17from sqlalchemy.orm.session import Session
18from sqlalchemy.pool import QueuePool
19from sqlalchemy.sql import and_, func, literal, or_
21from couchers.config import config
22from couchers.constants import DB_POOL_SIZE
23from couchers.models import (
24 Cluster,
25 ClusterRole,
26 ClusterSubscription,
27 FriendRelationship,
28 FriendStatus,
29 Geom,
30 Node,
31 TimezoneArea,
32 User,
33)
34from couchers.perf import register_perf_listeners
35from couchers.sql import where_users_column_visible
37if TYPE_CHECKING:
38 from couchers.context import CouchersContext
40# Register psycopg (psycopg3) as the default driver for postgresql:// URLs
41# This must happen before any engine is created
42registry.register("postgresql", "sqlalchemy.dialects.postgresql.psycopg", "PGDialect_psycopg")
44logger = logging.getLogger(__name__)
46tracer = trace.get_tracer(__name__)
49def apply_migrations() -> None:
50 alembic_dir = os.path.dirname(__file__) + "/../.."
51 cwd = os.getcwd()
52 try:
53 os.chdir(alembic_dir)
54 alembic_cfg = Config("alembic.ini")
55 # alembic screws up logging config by default, this tells it not to screw it up if being run at startup like this
56 alembic_cfg.set_main_option("dont_mess_up_logging", "False")
57 command.upgrade(alembic_cfg, "head")
58 finally:
59 os.chdir(cwd)
62@functools.cache
63def _get_base_engine() -> Engine:
64 engine = create_engine(
65 config.DATABASE_CONNECTION_STRING,
66 # checks that the connections in the pool are alive before using them, which avoids the "server closed the
67 # connection unexpectedly" errors
68 pool_pre_ping=True,
69 # one connection per thread
70 poolclass=QueuePool,
71 # each process keeps its own pool, so total connections ~= process count * pool_size, kept under postgres
72 # max_connections. ~2 per thread since a thread can hold two connections at once (handler + _store_log,
73 # or the jobs worker's own session_scope + the handler's).
74 pool_size=DB_POOL_SIZE,
75 max_overflow=0,
76 )
77 register_perf_listeners(engine)
78 return engine
81@contextmanager
82def session_scope() -> Generator[Session]:
83 with tracer.start_as_current_span("session_scope") as rollspan:
84 with Session(_get_base_engine()) as session:
85 session.begin()
86 try:
87 if logger.isEnabledFor(logging.DEBUG): 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true
88 try:
89 frame = inspect.stack()[2]
90 filename_line = f"{frame.filename}:{frame.lineno}"
91 except Exception as e:
92 filename_line = "{unknown file}"
93 backend_pid = session.execute(text("SELECT pg_backend_pid();")).scalar_one()
94 logger.debug(f"SScope: got {backend_pid=} at {filename_line}")
95 rollspan.set_attribute("db.backend_pid", backend_pid)
96 rollspan.set_attribute("db.filename_line", filename_line)
97 rollspan.set_attribute("rpc.thread", get_ident())
98 rollspan.set_attribute("rpc.pid", getpid())
100 yield session
101 session.commit()
102 except:
103 session.rollback()
104 raise
105 finally:
106 if logger.isEnabledFor(logging.DEBUG): 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true
107 logger.debug(f"SScope: closed {backend_pid=}")
110def db_post_fork() -> None:
111 """
112 Fix post-fork issues with sqlalchemy
113 """
114 # see https://docs.sqlalchemy.org/en/20/core/pooling.html#using-connection-pools-with-multiprocessing-or-os-fork
115 _get_base_engine().dispose(close=False)
118def are_friends(session: Session, context: CouchersContext, other_user: int) -> bool:
119 query = select(FriendRelationship)
120 query = where_users_column_visible(query, context, FriendRelationship.from_user_id)
121 query = where_users_column_visible(query, context, FriendRelationship.to_user_id)
122 query = query.where(
123 or_(
124 and_(FriendRelationship.from_user_id == context.user_id, FriendRelationship.to_user_id == other_user),
125 and_(FriendRelationship.from_user_id == other_user, FriendRelationship.to_user_id == context.user_id),
126 )
127 ).where(FriendRelationship.status == FriendStatus.accepted)
128 return session.execute(query).scalar_one_or_none() is not None
131def get_parent_node_at_location(session: Session, shape: WKBElement) -> Node | None:
132 """
133 Finds the smallest node containing the shape.
135 Shape can be any PostGIS geo object, e.g., output from create_coordinate
136 """
138 # Find the lowest Node (in the Node tree) that contains the shape. By construction of nodes, the area of a sub-node
139 # must always be less than its parent Node, so no need to actually traverse the tree!
140 return (
141 session.execute(
142 select(Node).where(func.ST_Contains(Node.geom, shape)).order_by(func.ST_Area(Node.geom)).limit(1)
143 )
144 .scalars()
145 .one_or_none()
146 )
149def _get_node_parents_recursive_cte_subquery(node_id: int) -> Subquery:
150 parents = (
151 select(Node.id, Node.parent_node_id, literal(0).label("level"))
152 .where(Node.id == node_id)
153 .cte("parents", recursive=True)
154 )
156 return select(
157 parents.union(
158 select(Node.id, Node.parent_node_id, (parents.c.level + 1).label("level")).join(
159 parents, Node.id == parents.c.parent_node_id
160 )
161 )
162 ).subquery()
165def get_node_parents_recursively(session: Session, node_id: int) -> Sequence[Row[tuple[int, int, int, Cluster]]]:
166 subquery = _get_node_parents_recursive_cte_subquery(node_id)
167 return session.execute(
168 select(subquery, Cluster)
169 .join(Cluster, Cluster.parent_node_id == subquery.c.id)
170 .where(Cluster.is_official_cluster)
171 .order_by(subquery.c.level.desc())
172 ).all()
175def _can_moderate_any_cluster(session: Session, user_id: int, cluster_ids: list[int]) -> bool:
176 query = select(
177 (
178 select(true())
179 .select_from(ClusterSubscription)
180 .where(ClusterSubscription.role == ClusterRole.admin)
181 .where(ClusterSubscription.user_id == user_id)
182 .where(ClusterSubscription.cluster_id.in_(cluster_ids))
183 ).exists()
184 )
185 return session.execute(query).scalar_one()
188def can_moderate_node(session: Session, user_id: int, node_id: int) -> bool:
189 """
190 Returns True if the user_id can moderate the given node (i.e., if they are admin of any community that is a parent of the node)
191 """
192 subquery = _get_node_parents_recursive_cte_subquery(node_id)
193 query = select(
194 (
195 select(true())
196 .select_from(ClusterSubscription)
197 .where(ClusterSubscription.role == ClusterRole.admin)
198 .where(ClusterSubscription.user_id == user_id)
199 .join(Cluster, Cluster.id == ClusterSubscription.cluster_id)
200 .where(Cluster.is_official_cluster)
201 .where(Cluster.parent_node_id == subquery.c.id)
202 ).exists()
203 )
204 return session.execute(query).scalar_one()
207def can_moderate_at(session: Session, user_id: int, shape: Geom) -> bool:
208 """
209 Returns True if the user_id can moderate a given geo-shape (i.e., if the shape is contained in any Node that the user is an admin of)
210 """
211 query = select(
212 (
213 select(true())
214 .select_from(ClusterSubscription)
215 .where(ClusterSubscription.role == ClusterRole.admin)
216 .where(ClusterSubscription.user_id == user_id)
217 .join(Cluster, Cluster.id == ClusterSubscription.cluster_id)
218 .join(Node, and_(Cluster.is_official_cluster, Node.id == Cluster.parent_node_id))
219 .where(func.ST_Contains(Node.geom, shape))
220 ).exists()
221 )
222 return session.execute(query).scalar_one()
225def is_user_in_node_geography(session: Session, user_id: int, node_id: int) -> bool:
226 """
227 Returns True if the user's location is geographically contained within the node's boundary.
228 This is used to check if a user can leave a community - users cannot leave communities
229 that contain their geographic location.
230 """
231 query = select(
232 (
233 select(true())
234 .select_from(User)
235 .join(Node, func.ST_Contains(Node.geom, User.geom))
236 .where(User.id == user_id)
237 .where(Node.id == node_id)
238 ).exists()
239 )
240 return session.execute(query).scalar_one()
243def timezone_at_coordinate(session: Session, geom: WKBElement) -> str | None:
244 tzid = session.execute(
245 select(TimezoneArea.tzid).where(func.ST_Contains(TimezoneArea.geom, geom))
246 ).scalar_one_or_none()
247 return tzid