sqlalchemy_oso_cloud.orm

 1"""
 2Utilities for [declaratively mapping](https://docs.sqlalchemy.org/en/20/orm/mapping_styles.html#orm-declarative-mapping)
 3[authorization data](https://www.osohq.com/docs/authorization-data) in your ORM models.
 4"""
 5
 6from typing import Any, Callable, Optional, Protocol
 7
 8from sqlalchemy.orm import MappedColumn, Relationship, mapped_column, relationship
 9from typing_extensions import Concatenate, ParamSpec, TypeVar
10
11
12class Resource:
13  """
14  A mixin to indicate that an ORM model corresponds to an Oso resource.
15  """
16  pass
17
18_RELATION_INFO_KEY = "_oso.relation"
19_ATTRIBUTE_INFO_KEY = "_oso.attribute"
20_REMOTE_RELATION_INFO_KEY = "_oso.remote_relation"
21
22P = ParamSpec('P')
23T = TypeVar('T')
24R = TypeVar('R', covariant=True)
25
26def _wrap(func: Callable[P, Any]) -> Callable[[Callable[P, T]], Callable[P, T]]:
27    """Wrap a SQLAlchemy function in a type-safe way."""
28    def decorator(wrapper: Callable[P, T]) -> Callable[P, T]:
29      def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
30          return wrapper(*args, **kwargs)
31      return wrapped
32    return decorator
33
34class _WithExtraKwargs(Protocol[P, R]):
35    def __call__(self, remote_resource_name: str, remote_relation_key: Optional[str] = None, *args: P.args, **kwargs: P.kwargs) -> R:
36        ...
37
38def _add_params(wrapped_func: Callable[P, R]) -> Callable[[_WithExtraKwargs[P, R]], _WithExtraKwargs[P, R]]:
39  """Adds extra keyword parameters to `remote_relation` in order to support static type checking."""
40  def decorator(wrapper: Callable[Concatenate[str, Optional[str], P], R]) -> _WithExtraKwargs[P, R]:
41    def wrapped(remote_resource_name: str, remote_relation_key: Optional[str] = None, *args: P.args, **kwargs: P.kwargs) -> R:
42      return wrapper(remote_resource_name, remote_relation_key, *args, **kwargs)
43    return wrapped
44  return decorator
45
46@_wrap(relationship)
47def relation(*args, **kwargs) -> Relationship:
48  """
49  A wrapper around [`sqlalchemy.orm.relationship`](https://docs.sqlalchemy.org/en/20/orm/relationship_api.html#sqlalchemy.orm.relationship)
50  that indicates that the relationship corresponds to `has_relation` facts in Oso with the following three arguments:
51  1. the resource this relationship is declared on,
52  2. the name of this relationship, and
53  3. the resource that the relationship points to.
54
55  Accepts all of the same arguments as [`sqlalchemy.orm.relationship`](https://docs.sqlalchemy.org/en/20/orm/relationship_api.html#sqlalchemy.orm.relationship).
56  """
57  rel = relationship(*args, **kwargs)
58  rel.info[_RELATION_INFO_KEY] = None
59  return rel
60
61@_wrap(mapped_column)
62def attribute(*args, **kwargs) -> MappedColumn:
63  """
64  A wrapper around [`sqlalchemy.orm.mapped_column`](https://docs.sqlalchemy.org/en/20/orm/mapping_api.html#sqlalchemy.orm.mapped_column)
65  that indicates that the attribute corresponds to `has_{attribute_name}` facts in Oso with the following two arguments:
66  1. the resource this attribute is declared on, and
67  2. the attribute value.
68
69  Accepts all of the same arguments as [`sqlalchemy.orm.mapped_column`](https://docs.sqlalchemy.org/en/20/orm/mapping_api.html#sqlalchemy.orm.mapped_column).
70  """
71  col = mapped_column(*args, **kwargs)
72  col.column.info[_ATTRIBUTE_INFO_KEY] = None
73  return col
74
75@_add_params(mapped_column)
76def remote_relation(remote_resource_name: str, remote_relation_key: Optional[str] = None, *args, **kwargs) -> MappedColumn:
77  """
78  A wrapper around [`sqlalchemy.orm.mapped_column`](https://docs.sqlalchemy.org/en/20/orm/mapping_api.html#sqlalchemy.orm.mapped_column)
79  that indicates that the attribute corresponds to `has_relation` facts (to a resource not defined in the local database) in Oso with the following three arguments:
80  1. the resource this attribute is declared on, and
81  2. the name of this relationship, and
82  3. the resource that the relationship points to.
83
84  Note: this is not a [`sqlalchemy.orm.relationship`](https://docs.sqlalchemy.org/en/20/orm/relationship_api.html#sqlalchemy.orm.relationship).
85
86  Accepts all of the same arguments as [`sqlalchemy.orm.mapped_column`](https://docs.sqlalchemy.org/en/20/orm/mapping_api.html#sqlalchemy.orm.mapped_column).
87  Also accepts the following additional arguments:
88  :param remote_resource_name: the name of the remote resource
89  :param remote_relation_key: (optional) the name of the relation on the remote resource. If not provided, the name of the relation will be inferred from the name of the column.
90  """
91  col = mapped_column(*args, **kwargs)
92  col.column.info[_REMOTE_RELATION_INFO_KEY] = (remote_resource_name, remote_relation_key)
93  return col