authentik.enterprise.search.fields

  1"""DjangoQL search"""
  2
  3from collections import OrderedDict, defaultdict
  4from collections.abc import Generator
  5
  6from django.db import connection
  7from django.db.models import Model, Q
  8from djangoql.compat import text_type
  9from djangoql.schema import StrField
 10from djangoql.serializers import DjangoQLSchemaSerializer
 11
 12
 13class JSONSearchField(StrField):
 14    """JSON field for DjangoQL"""
 15
 16    model: Model
 17
 18    def __init__(
 19        self,
 20        model=None,
 21        name=None,
 22        nullable=None,
 23        suggest_nested=False,
 24        fixed_structure: OrderedDict | None = None,
 25    ):
 26        # Set this in the constructor to not clobber the type variable
 27        self.type = "relation"
 28        self.suggest_nested = suggest_nested
 29        self.fixed_structure = fixed_structure
 30        super().__init__(model, name, nullable)
 31
 32    def get_lookup(self, path, operator, value):
 33        search = "__".join(path)
 34        op, invert = self.get_operator(operator)
 35        q = Q(**{f"{search}{op}": self.get_lookup_value(value)})
 36        return ~q if invert else q
 37
 38    def json_field_keys(self) -> Generator[tuple[str]]:
 39        with connection.cursor() as cursor:
 40            cursor.execute(f"""
 41                WITH RECURSIVE "{self.name}_keys" AS (
 42                    SELECT
 43                        ARRAY[jsonb_object_keys("{self.name}")] AS key_path_array,
 44                        "{self.name}" -> jsonb_object_keys("{self.name}") AS value
 45                    FROM {self.model._meta.db_table}
 46                    WHERE "{self.name}" IS NOT NULL
 47                        AND jsonb_typeof("{self.name}") = 'object'
 48
 49                    UNION ALL
 50
 51                    SELECT
 52                        ck.key_path_array || jsonb_object_keys(ck.value),
 53                        ck.value -> jsonb_object_keys(ck.value) AS value
 54                    FROM "{self.name}_keys" ck
 55                    WHERE jsonb_typeof(ck.value) = 'object'
 56                ),
 57
 58                unique_paths AS (
 59                    SELECT DISTINCT key_path_array
 60                    FROM "{self.name}_keys"
 61                )
 62
 63                SELECT key_path_array FROM unique_paths;
 64            """)  # nosec
 65            return (x[0] for x in cursor.fetchall())
 66
 67    def get_fixed_structure(self, serializer: DjangoQLSchemaSerializer) -> OrderedDict:
 68        new_dict = OrderedDict()
 69        if not self.fixed_structure:
 70            return new_dict
 71        new_dict.setdefault(self.relation(), {})
 72        for key, value in self.fixed_structure.items():
 73            new_dict[self.relation()][key] = serializer.serialize_field(value)
 74            if isinstance(value, JSONSearchField):
 75                new_dict.update(value.get_nested_options(serializer))
 76        return new_dict
 77
 78    def get_nested_options(self, serializer: DjangoQLSchemaSerializer) -> OrderedDict:
 79        """Get keys of all nested objects to show autocomplete"""
 80        if not self.suggest_nested:
 81            if self.fixed_structure:
 82                return self.get_fixed_structure(serializer)
 83            return OrderedDict()
 84
 85        def recursive_function(parts: list[str], parent_parts: list[str] | None = None):
 86            if not parent_parts:
 87                parent_parts = []
 88            path = parts.pop(0)
 89            parent_parts.append(path)
 90            relation_key = "_".join(parent_parts)
 91            if len(parts) > 1:
 92                out_dict = {
 93                    relation_key: {
 94                        parts[0]: {
 95                            "type": "relation",
 96                            "relation": f"{relation_key}_{parts[0]}",
 97                        }
 98                    }
 99                }
100                child_paths = recursive_function(parts.copy(), parent_parts.copy())
101                child_paths.update(out_dict)
102                return child_paths
103            else:
104                return {relation_key: {parts[0]: {}}}
105
106        relation_structure = defaultdict(dict)
107
108        for relations in self.json_field_keys():
109            result = recursive_function([self.relation()] + relations)
110            for relation_key, value in result.items():
111                for sub_relation_key, sub_value in value.items():
112                    if not relation_structure[relation_key].get(sub_relation_key, None):
113                        relation_structure[relation_key][sub_relation_key] = sub_value
114                    else:
115                        relation_structure[relation_key][sub_relation_key].update(sub_value)
116
117        final_dict = defaultdict(dict)
118
119        for key, value in relation_structure.items():
120            for sub_key, sub_value in value.items():
121                if not sub_value:
122                    final_dict[key][sub_key] = {
123                        "type": "str",
124                        "nullable": True,
125                    }
126                else:
127                    final_dict[key][sub_key] = sub_value
128        return OrderedDict(final_dict)
129
130    def relation(self) -> str:
131        return f"{self.model._meta.app_label}.{self.model._meta.model_name}_{self.name}"
132
133
134class ChoiceSearchField(StrField):
135    def __init__(self, model=None, name=None, nullable=None):
136        super().__init__(model, name, nullable, suggest_options=True)
137
138    def get_options(self, search):
139        result = []
140        choices = self._field_choices()
141        if choices:
142            search = search.lower()
143            for c in choices:
144                choice = text_type(c[0])
145                if search in choice.lower():
146                    result.append(choice)
147        return result