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