sqlalchemy_oso_cloud

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

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

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

146def get_oso() -> Oso:
147  """
148  Get the Oso Cloud client that was created with `init`.
149
150  :return: The Oso Cloud client.
151  """
152  global oso
153  if oso is None:
154    raise RuntimeError("sqlalchemy_oso_cloud must be initialized before getting the Oso client")
155  return oso

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

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

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

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