-
Notifications
You must be signed in to change notification settings - Fork 116
/
fields.py
115 lines (87 loc) · 4.06 KB
/
fields.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import logging
import django
from django.db import models
from django.db.models.sql.where import WhereNode
from django.conf import settings
from .exceptions import EmptyTenant
from .utils import get_current_tenant, get_tenant_column, get_tenant_filters
logger = logging.getLogger(__name__)
class TenantForeignKey(models.ForeignKey):
"""
Should be used in place of models.ForeignKey for all foreign key relationships to
subclasses of TenantModel.
Adds additional clause to JOINs over this relation to include tenant_id in the JOIN
on the TenantModel.
Adds clause to forward accesses through this field to include tenant_id in the
TenantModel lookup.
"""
# Override
def get_extra_descriptor_filter(self, instance):
"""
Return an extra filter condition for related object fetching when
user does 'instance.fieldname', that is the extra filter is used in
the descriptor of the field.
The filter should be either a dict usable in .filter(**kwargs) call or
a Q-object. The condition will be ANDed together with the relation's
joining columns.
A parallel method is get_extra_restriction() which is used in
JOIN and subquery conditions.
"""
current_tenant = get_current_tenant()
if current_tenant:
return get_tenant_filters(self.related_model)
empty_tenant_message = (
f"TenantForeignKey field {self.model.__name__}.{self.name} "
"accessed without a current tenant set. "
"This may cause issues in a partitioned environment. "
"Recommend calling set_current_tenant() before accessing "
"this field."
)
if getattr(settings, "TENANT_STRICT_MODE", False):
raise EmptyTenant(empty_tenant_message)
logger.warning(empty_tenant_message)
return super().get_extra_descriptor_filter(instance)
# Override
# Django 4.0 removed the where_class argument from this method, so
# depending on the version we define the function with a different
# signature.
# pylint: disable=unused-argument,arguments-differ
if django.VERSION >= (4, 0):
def get_extra_restriction(self, alias, related_alias):
return self.get_extra_restriction_citus(alias, related_alias)
else:
def get_extra_restriction(self, where_class, alias, related_alias):
return self.get_extra_restriction_citus(alias, related_alias)
def get_extra_restriction_citus(self, alias, related_alias):
"""
Return a pair condition used for joining and subquery pushdown. The
condition is something that responds to as_sql(compiler, connection)
method.
Note that currently referring both the 'alias' and 'related_alias'
will not work in some conditions, like subquery pushdown.
A parallel method is get_extra_descriptor_filter() which is used in
instance.fieldname related object fetching.
"""
if not (related_alias and alias):
return None
# Fetch tenant column names for both sides of the relation
lhs_model = self.model
rhs_model = self.related_model
lhs_tenant_id = get_tenant_column(lhs_model)
rhs_tenant_id = get_tenant_column(rhs_model)
# Fetch tenant fields for both sides of the relation
lhs_tenant_field = lhs_model._meta.get_field(lhs_tenant_id)
rhs_tenant_field = rhs_model._meta.get_field(rhs_tenant_id)
# Get references to both tenant columns
lookup_lhs = lhs_tenant_field.get_col(related_alias)
lookup_rhs = rhs_tenant_field.get_col(alias)
# Create "AND lhs.tenant_id = rhs.tenant_id" as a new condition
lookup = lhs_tenant_field.get_lookup("exact")(lookup_lhs, lookup_rhs)
condition = WhereNode()
condition.add(lookup, "AND")
return condition
class TenantOneToOneField(models.OneToOneField, TenantForeignKey):
# Override
def __init__(self, *args, **kwargs):
kwargs["unique"] = False
super().__init__(*args, **kwargs)