github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/clients/python-wrapper/lakefs/reference.py (about)

     1  """
     2  Module containing lakeFS reference implementation
     3  """
     4  
     5  from __future__ import annotations
     6  
     7  from typing import Optional, Generator, Union
     8  
     9  import lakefs_sdk
    10  
    11  from lakefs.models import Commit, Change, CommonPrefix, ObjectInfo, _OBJECT
    12  from lakefs.client import Client, _BaseLakeFSObject
    13  from lakefs.exceptions import api_exception_handler
    14  from lakefs.object import StoredObject
    15  
    16  
    17  class Reference(_BaseLakeFSObject):
    18      """
    19      Class representing a reference in lakeFS.
    20      """
    21      _repo_id: str
    22      _id: str
    23      _commit: Optional[Commit] = None
    24  
    25      def __init__(self, repository_id: str, reference_id: str, client: Optional[Client] = None) -> None:
    26          """Return a reference to a lakeFS commit.
    27  
    28          :param repository_id: the repository holding the commit
    29          :param reference_id: a reference expression to the commit
    30  
    31          Any reference expression can be used as a reference_id, for example:
    32  
    33          - 'main' (head of 'main' branch)
    34          - 'main@' (head of 'main' branch, only committed objects)
    35          - 'my_tag~3' (3 commits before 'my_tag')
    36  
    37          See https://docs.lakefs.io/understand/model.html#ref-expressions for
    38          details.
    39          """
    40          self._repo_id = repository_id
    41          self._id = reference_id
    42          super().__init__(client)
    43  
    44      @property
    45      def repo_id(self) -> str:
    46          """
    47          Return the repository id for this reference
    48          """
    49          return self._repo_id
    50  
    51      @property
    52      def id(self) -> str:
    53          """
    54          Returns the reference id
    55          """
    56          return self._id
    57  
    58      def objects(self,
    59                  max_amount: Optional[int] = None,
    60                  after: Optional[str] = None,
    61                  prefix: Optional[str] = None,
    62                  delimiter: Optional[str] = None,
    63                  **kwargs) -> Generator[StoredObject | CommonPrefix]:
    64          """
    65          Returns an object generator for this reference, the generator can yield either a StoredObject or a CommonPrefix
    66          object depending on the listing parameters provided.
    67  
    68          :param max_amount: Stop showing changes after this amount
    69          :param after: Return items after this value
    70          :param prefix: Return items prefixed with this value
    71          :param delimiter: Group common prefixes by this delimiter
    72          :param kwargs: Additional Keyword Arguments to send to the server
    73          :raise NotFoundException: if this reference or other_ref does not exist
    74          :raise NotAuthorizedException: if user is not authorized to perform this operation
    75          :raise ServerException: for any other errors
    76          """
    77  
    78          for res in generate_listing(self._client.sdk_client.objects_api.list_objects,
    79                                      repository=self._repo_id,
    80                                      ref=self._id,
    81                                      max_amount=max_amount,
    82                                      after=after,
    83                                      prefix=prefix,
    84                                      delimiter=delimiter,
    85                                      **kwargs):
    86              type_class = ObjectInfo if res.path_type == _OBJECT else CommonPrefix
    87              yield type_class(**res.dict())
    88  
    89      def log(self, max_amount: Optional[int] = None, **kwargs) -> Generator[Commit]:
    90          """
    91          Returns a generator of commits starting with this reference id
    92  
    93          :param max_amount: (Optional) limits the amount of results to return from the server
    94          :param kwargs: Additional Keyword Arguments to send to the server
    95          :raise NotFoundException: if reference by this id does not exist
    96          :raise NotAuthorizedException: if user is not authorized to perform this operation
    97          :raise ServerException: for any other errors
    98          """
    99          for res in generate_listing(self._client.sdk_client.refs_api.log_commits, self._repo_id, self._id,
   100                                      max_amount=max_amount, **kwargs):
   101              yield Commit(**res.dict())
   102  
   103      def get_commit(self) -> Commit:
   104          """
   105          Returns the underlying commit referenced by this reference id
   106  
   107          :raise NotFoundException: if this reference does not exist
   108          :raise NotAuthorizedException: if user is not authorized to perform this operation
   109          :raise ServerException: for any other errors
   110          """
   111          if self._commit is None:
   112              with api_exception_handler():
   113                  commit = self._client.sdk_client.commits_api.get_commit(self._repo_id, self._id)
   114                  self._commit = Commit(**commit.dict())
   115          return self._commit
   116  
   117      def diff(self,
   118               other_ref: ReferenceType,
   119               max_amount: Optional[int] = None,
   120               after: Optional[str] = None,
   121               prefix: Optional[str] = None,
   122               delimiter: Optional[str] = None,
   123               **kwargs) -> Generator[Change]:
   124          """
   125          Returns a diff generator of changes between this reference and other_ref
   126  
   127          :param other_ref: The other ref to diff against
   128          :param max_amount: Stop showing changes after this amount
   129          :param after: Return items after this value
   130          :param prefix: Return items prefixed with this value
   131          :param delimiter: Group common prefixes by this delimiter
   132          :param kwargs: Additional Keyword Arguments to send to the server
   133          :raise NotFoundException: if this reference or other_ref does not exist
   134          :raise NotAuthorizedException: if user is not authorized to perform this operation
   135          :raise ServerException: for any other errors
   136          """
   137          other_ref_id = other_ref if isinstance(other_ref, str) else other_ref.id
   138          for diff in generate_listing(self._client.sdk_client.refs_api.diff_refs,
   139                                       repository=self._repo_id,
   140                                       left_ref=self._id,
   141                                       right_ref=other_ref_id,
   142                                       after=after,
   143                                       max_amount=max_amount,
   144                                       prefix=prefix,
   145                                       delimiter=delimiter,
   146                                       **kwargs):
   147              yield Change(**diff.dict())
   148  
   149      def merge_into(self, destination_branch: ReferenceType, **kwargs) -> str:
   150          """
   151          Merge this reference into destination branch
   152  
   153          :param destination_branch: The merge destination (either ID or branch object)
   154          :param kwargs: Additional Keyword Arguments to send to the server
   155          :return: The reference id of the merge commit
   156          :raise NotFoundException: if reference by this id does not exist, or branch doesn't exist
   157          :raise NotAuthorizedException: if user is not authorized to perform this operation
   158          :raise ServerException: for any other errors
   159          """
   160          branch_id = destination_branch if isinstance(destination_branch, str) else destination_branch.id
   161          with api_exception_handler():
   162              merge = lakefs_sdk.Merge(**kwargs)
   163              res = self._client.sdk_client.refs_api.merge_into_branch(self._repo_id,
   164                                                                       self._id,
   165                                                                       branch_id,
   166                                                                       merge=merge)
   167              return res.reference
   168  
   169      def object(self, path: str) -> StoredObject:  # pylint: disable=C0103
   170          """
   171          Returns an Object class representing a lakeFS object with this repo id, reference id and path
   172  
   173          :param path: The object's path
   174          """
   175          return StoredObject(self._repo_id, self._id, path, self._client)
   176  
   177      def __repr__(self):
   178          class_name = self.__class__.__name__
   179          return f'{class_name}(repository="{self.repo_id}", id="{self.id}")'
   180  
   181  
   182  def generate_listing(func, *args, max_amount: Optional[int] = None, **kwargs):
   183      """
   184      Generic generator function, for lakefs-sdk listings functionality
   185  
   186      :param func: The listing function
   187      :param args: The function args
   188      :param max_amount: The max amount of objects to generate
   189      :param kwargs: The function kwargs
   190      :return: A generator based on the listing function
   191      """
   192      has_more = True
   193      with api_exception_handler():
   194          while has_more:
   195              page = func(*args, **kwargs)
   196              has_more = page.pagination.has_more
   197              kwargs["after"] = page.pagination.next_offset
   198              for res in page.results:
   199                  yield res
   200  
   201                  if max_amount is not None:
   202                      max_amount -= 1
   203                      if max_amount <= 0:
   204                          return
   205  
   206  
   207  ReferenceType = Union[str, Reference, Commit]