python/samba/domain/models/query.py

   1 # Unix SMB/CIFS implementation.
   2 #
   3 # Query class for the ORM to the Ldb database.
   4 #
   5 # Copyright (C) Catalyst.Net Ltd. 2023
   6 #
   7 # Written by Rob van der Linde <rob@catalyst.net.nz>
   8 #
   9 # This program is free software; you can redistribute it and/or modify
  10 # it under the terms of the GNU General Public License as published by
  11 # the Free Software Foundation; either version 3 of the License, or
  12 # (at your option) any later version.
  13 #
  14 # This program is distributed in the hope that it will be useful,
  15 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  16 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  17 # GNU General Public License for more details.
  18 #
  19 # You should have received a copy of the GNU General Public License
  20 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  21 #
  22
  23 import re
  24
  25 from .exceptions import NotFound, MultipleObjectsReturned
  26 from .registry import MODELS
  27
  28 RE_SPLIT_CAMELCASE = re.compile(r"[A-Z](?:[a-z]+|[A-Z]*(?=[A-Z]|$))")
  29
  30
  31 class Query:
  32     """Simple Query class used by the `Model.query` method."""
  33
  34     def __init__(self, model, samdb, result, polymorphic):
  35         self.model = model
  36         self.samdb = samdb
  37         self.result = result
  38         self.count = result.count
  39         self.name = " ".join(RE_SPLIT_CAMELCASE.findall(model.__name__)).lower()
  40         self.polymorphic = polymorphic
  41
  42     def __iter__(self):
  43         """Loop over Query class yields Model instances."""
  44         for message in self.result:
  45             yield self._from_message(message)
  46
  47     def __repr__(self):
  48         """Provide repr method that provides more useful output in the shell."""
  49         return f"<Query {list(self)}>"
  50
  51     def _from_message(self, message):
  52         """Returns the model class to use to construct instances.
  53
  54         If polymorphic query is enabled it will use the last item from
  55         the objectClass list.
  56
  57         Otherwise, it will use the model from the queryset.
  58         """
  59         if self.polymorphic:
  60             object_class = str(message["objectClass"][-1])
  61             model = MODELS.get(object_class, self.model)
  62         else:
  63             model = self.model
  64
  65         return model._from_message(self.samdb, message)
  66
  67     def first(self):
  68         """Returns the first item in the Query or None for no results."""
  69         if self.count:
  70             return self._from_message(self.result[0])
  71
  72     def last(self):
  73         """Returns the last item in the Query or None for no results."""
  74         if self.count:
  75             return self._from_message(self.result[-1])
  76
  77     def get(self):
  78         """Returns one item or None if no results were found.
  79
  80         :returns: Model instance or None if not found.
  81         :raises MultipleObjectsReturned: if more than one results were returned
  82         """
  83         if self.count > 1:
  84             raise MultipleObjectsReturned(
  85                 f"More than one {self.name} objects returned (got {self.count}).")
  86         elif self.count:
  87             return self._from_message(self.result[0])
  88
  89     def one(self):
  90         """Must return EXACTLY one item or raise an exception.
  91
  92         :returns: Model instance
  93         :raises NotFound: if no results were returned
  94         :raises MultipleObjectsReturned: if more than one results were returned
  95         """
  96         if self.count < 1:
  97             raise NotFound(
  98                 f"{self.name.capitalize()} matching query not found")
  99         elif self.count > 1:
 100             raise MultipleObjectsReturned(
 101                 f"More than one {self.name} objects returned (got {self.count}).")
 102         else:
 103             return self._from_message(self.result[0])