sqlalchemy_oso_cloud

16class Session(sqlalchemy.orm.Session):
17  """
18  A convenience wrapper around SQLAlchemy's built-in
19  [`sqlalchemy.orm.Session`](https://docs.sqlalchemy.org/orm/session_api.html#sqlalchemy%2Eorm%2ESession)
20  class.
21  
22  This class extends SQLAlchemy's Session to automatically use our custom `.Query` class
23  instead of the default [`sqlalchemy.orm.Query`](https://docs.sqlalchemy.org/orm/queryguide/query.html#sqlalchemy%2Eorm%2EQuery) class.
24  This is only useful if you intend to use SQLAlchemy's [legacy Query API](https://docs.sqlalchemy.org/orm/queryguide/query.html).
25  """
26  _query_cls: Type[Query] = Query
27
28  def __init__(self, *args, **kwargs):
29    """
30    Initialize a SQLAlchemy session with the `.Query` class extended to support authorization.
31    Accepts all of the same arguments as [`sqlalchemy.orm.Session`](https://docs.sqlalchemy.org/orm/session_api.html#sqlalchemy%2Eorm%2ESession),
32    except for `query_cls`.
33    """
34    if "query_cls" in kwargs:
35      raise ValueError("sqlalchemy_oso_cloud does not currently support combining with other query classes")
36    super().__init__(*args, **{ **kwargs, "query_cls": Query })
37
38  # Single entity overload
39  @overload # type: ignore[override]
40  def query(self, entity: Type[T], /) -> Query[T]: ...
41
42  # Single column overload
43  @overload
44  def query(self, column: InstrumentedAttribute[T], /) -> Query[Row[Tuple[T]]]: ...
45
46   # Two arguments - any combination of entities/columns
47  @overload
48  def query(self, 
49           arg1: Union[Type[T1], InstrumentedAttribute[T1]], 
50           arg2: Union[Type[T2], InstrumentedAttribute[T2]], /) -> Query[Row[Tuple[T1, T2]]]: ...
51
52  # Three arguments - any combination of entities/columns
53  @overload
54  def query(self, 
55           arg1: Union[Type[T1], InstrumentedAttribute[T1]], 
56           arg2: Union[Type[T2], InstrumentedAttribute[T2]], 
57           arg3: Union[Type[T3], InstrumentedAttribute[T3]], /) -> Query[Row[Tuple[T1, T2, T3]]]: ...
58
59  # Four arguments - any combination of entities/columns
60  @overload
61  def query(self, 
62           arg1: Union[Type[T1], InstrumentedAttribute[T1]], 
63           arg2: Union[Type[T2], InstrumentedAttribute[T2]], 
64           arg3: Union[Type[T3], InstrumentedAttribute[T3]], 
65           arg4: Union[Type[T4], InstrumentedAttribute[T4]], /) -> Query[Row[Tuple[T1, T2, T3, T4]]]: ...
66
67
68  # Fallback overload
69  @overload
70  def query(self, *entities: Any) -> Query[Any]: ...
71
72  def query(self, *entities, **kwargs) -> Query[Any]:
73      """
74      Returns a SQLAlchemy query extended to support authorization.
75
76      Returns a SQLAlchemy query extended to support authorization.
77      Accepts all of the same arguments as
78      [`sqlalchemy.orm.Session.query`](https://docs.sqlalchemy.org/en/20/orm/session_api.html#sqlalchemy%2Eorm%2ESession%2Equery).
79    
80      Single entity queries return Query[T].
81      Multi-entity and column queries return Query[Row[Tuple[...]]].
82      All other queries types return Query[Any].
83      """
84      return super().query(*entities, **kwargs)

14class Query(sqlalchemy.orm.Query[T]):
15  """
16  An extension of [`sqlalchemy.orm.Query`](https://docs.sqlalchemy.org/orm/queryguide/query.html#sqlalchemy%2Eorm%2EQuery)
17  that adds support for authorization.
18  """
19  
20  def __init__(self, *args, **kwargs):
21      super().__init__(*args, **kwargs)
22      self.oso = get_oso()
23
24  def authorized(self: Self, actor: Value, action: str, model: Optional[Type] = None ) -> Self:
25    """
26    Filter the query to only include resources that the given actor is authorized to perform the given action on.
27
28    :param actor: The actor performing the action.
29    :param action: The action the actor is performing.
30
31    :return: A new query that includes only the resources that the actor is authorized to perform the action on.
32    """
33    return _apply_authorization_options(self, actor, action, model)

125def init(registry: registry, **kwargs):
126  """
127  Initialize an Oso Cloud client configured to resolve authorization data from your
128  database as specified in your ORM models.
129  See `.orm` for more information on how to map your authorization data.
130
131  :param registry: The SQLAlchemy registry containing your models. For example, `Base.registry`.
132  :param kwargs: Additional keyword arguments to pass to the Oso client constructor, such as `url` and `api_key`.
133  """
134  global oso
135  if oso is not None:
136    raise RuntimeError("sqlalchemy_oso_cloud has already been initialized")
137  kwargs = { **kwargs }
138  if "url" not in kwargs:
139    kwargs["url"] = os.getenv("OSO_URL", "https://api.osohq.com")
140  if "api_key" not in kwargs:
141    kwargs["api_key"] = os.getenv("OSO_AUTH")
142  if "data_bindings" in kwargs:
143    # just need to conditionally close/delete the temporary file if it was created
144    raise NotImplementedError("manual data_bindings are not supported yet")
145  with NamedTemporaryFile(mode="w") as f:
146    config = generate_local_authorization_config(registry)
147    yaml.dump(config, f)
148    f.flush()
149    kwargs["data_bindings"] = f.name
150    oso = Oso(**kwargs)

152def get_oso() -> Oso:
153  """
154  Get the Oso Cloud client that was created with `init`.
155
156  :return: The Oso Cloud client.
157  """
158  global oso
159  if oso is None:
160    raise RuntimeError("sqlalchemy_oso_cloud must be initialized before getting the Oso client")
161  return oso

11class Select(sqlalchemy.sql.Select):
12    """A Select subclass that adds authorization functionality"""
13
14    inherit_cache = True
15    """Internal SQLAlchemy caching optimization"""
16    
17    
18    def __init__(self, *args, **kwargs):
19        super().__init__(*args, **kwargs)
20    
21    def authorized(self: Self, actor: Value, action: str) -> Self:
22        """Add authorization filtering to the select statement"""
23        return _apply_authorization_options(self, actor, action)

26def select(*args, **kwargs) -> Select:
27    """
28    Create an sqlalchemy_oso_cloud.Select() object
29
30    This is a drop-in replacement for sqlalchemy.select() that adds
31    authorization capabilities via the .authorized() method.
32    
33    Example:
34        from sqlalchemy_oso_cloud import select
35
36        stmt = select(Document).where(Document.private == True).authorized(actor, "read")
37        documents = session.execute(stmt)
38    """
39    return Select(*args, **kwargs)

43def authorized(actor: Value, action: str, model: Type) -> LoaderCriteriaOption:
44    """
45    Create authorization options for use with .options()
46    
47    This function can be used with both Select and Query objects from the SQLAlchemy library:
48    
49    Examples:
50        # With standard SQLAlchemy select
51        from sqlalchemy import select
52        from sqlalchemy_oso_cloud import authorized
53
54        stmt = select(Document).options(*authorized(user, "read", Document))
55        docs = session.execute(stmt).scalars().all()
56
57        # With Query
58        docs = session.query(Document).options(*authorized(user, "read", Document)).all()
59
60    :param actor: The actor performing the action
61    :param action: The action the actor is performing  
62    :param models: The model classes to authorize against
63    :return: List of loader criteria options for use with .options()
64    """
65    
66    if not model:
67        raise ValueError("Must provide a model to authorize against.")
68   
69    if not issubclass(model, Resource):
70        raise ValueError(f"Model {model.__name__} must inherit from Resource to use authorization")
71
72    
73    auth_criteria = create_auth_criteria_for_model(model, actor, action)
74
75    return with_loader_criteria(
76            model,
77            auth_criteria,
78            include_aliases=True
79        )

105def _apply_authorization_options(query_obj: Union["Query",  "Select"], actor: Value, action: str, model: Optional[Type] = None):
106    """
107    Apply authorization to any query-like object that has column_descriptions and options()
108    
109    This works with both Select and Query objects.
110    """
111
112    if model is not None:
113        auth_option = authorized(actor, action, model)
114        return query_obj.options(auth_option)
115    else:
116        auth_options = _authorize_all_models(query_obj, actor, action)
117        return query_obj.options(*auth_options)