pgalchemy 0.1.2__tar.gz

pgalchemy-0.1.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Matthew Beatty
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pgalchemy-0.1.2/PKG-INFO

@@ -0,0 +1,83 @@
+Metadata-Version: 2.1
+Name: pgalchemy
+Version: 0.1.2
+Summary: 
+Author: Matthew Beatty
+Author-email: beattyml1@gmail.com
+Requires-Python: >=3.11,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: alembic (>=1.13.2,<2.0.0)
+Requires-Dist: alembic-utils (>=0.8.4,<0.9.0)
+Requires-Dist: sqlalchemy (>=2.0.32,<3.0.0)
+Description-Content-Type: text/markdown
+
+# pg-rls-sqlalchemy
+
+Work in progress. 
+
+SQLAlchemy and Alembic support for Postgres features like:
+- Row Level Security (RLS)
+- Policies
+
+Built on top of alembic_utils but provides a more usable interface and a few missing features
+
+## Installation
+
+```shell
+pip install pg-rls-sqlalchemy
+```
+
+OR 
+
+```shell
+poetry add pg-rls-sqlalchemy
+```
+
+## Usage
+
+### Using RLS BaseModel
+Recommended most projects. This is for projects with majority of tables using RLS which will also be almost all new projects using this library.
+
+```python
+
+from sqlalchemy.orm import declarative_base
+from pg_rls import rls_base, policy, Policy, PolicyType, PolicyCommands
+
+BaseModel = rls_base(declarative_base())
+
+
+@policy(Policy("pol_my_models_select_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.SELECT, using="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_delete_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.DELETE, using="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_update_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.UPDATE, using="user_id == auth.uid()", with_check="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_update_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.INSERT, with_check="user_id == auth.uid()"))
+# Equivalent to:
+# @policy(Policy("pol_my_models_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.ALL, using="user_id == auth.uid()", with_check="user_id == auth.uid()"))
+class MyModel(BaseModel):
+    ...
+```
+
+### Using RLS Decorator
+Only intended for projects with majority of tables without RLS enabled. Usually only for existing projects with most tables not protected using RLS that are only using RLS for a niche use case
+
+This is not recommended for other use cases as it makes it easy for a developer to forget to enable RLS and expose a security vulnerability.
+```python
+
+from sqlalchemy.orm import declarative_base
+from pg_rls import rls, policy, Policy, PolicyType, PolicyCommands
+
+BaseModel = declarative_base()
+
+@rls()
+@policy(Policy("pol_my_models_select_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.SELECT, using="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_delete_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.DELETE, using="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_update_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.UPDATE, using="user_id == auth.uid()", with_check="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_update_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.INSERT, with_check="user_id == auth.uid()"))
+# Equivalent to:
+# @policy(Policy("pol_my_models_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.ALL, using="user_id == auth.uid()", with_check="user_id == auth.uid()"))
+class MyModel(BaseModel):
+    ...
+```
+
+

pgalchemy-0.1.2/README.md

@@ -0,0 +1,67 @@
+# pg-rls-sqlalchemy
+
+Work in progress. 
+
+SQLAlchemy and Alembic support for Postgres features like:
+- Row Level Security (RLS)
+- Policies
+
+Built on top of alembic_utils but provides a more usable interface and a few missing features
+
+## Installation
+
+```shell
+pip install pg-rls-sqlalchemy
+```
+
+OR 
+
+```shell
+poetry add pg-rls-sqlalchemy
+```
+
+## Usage
+
+### Using RLS BaseModel
+Recommended most projects. This is for projects with majority of tables using RLS which will also be almost all new projects using this library.
+
+```python
+
+from sqlalchemy.orm import declarative_base
+from pg_rls import rls_base, policy, Policy, PolicyType, PolicyCommands
+
+BaseModel = rls_base(declarative_base())
+
+
+@policy(Policy("pol_my_models_select_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.SELECT, using="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_delete_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.DELETE, using="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_update_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.UPDATE, using="user_id == auth.uid()", with_check="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_update_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.INSERT, with_check="user_id == auth.uid()"))
+# Equivalent to:
+# @policy(Policy("pol_my_models_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.ALL, using="user_id == auth.uid()", with_check="user_id == auth.uid()"))
+class MyModel(BaseModel):
+    ...
+```
+
+### Using RLS Decorator
+Only intended for projects with majority of tables without RLS enabled. Usually only for existing projects with most tables not protected using RLS that are only using RLS for a niche use case
+
+This is not recommended for other use cases as it makes it easy for a developer to forget to enable RLS and expose a security vulnerability.
+```python
+
+from sqlalchemy.orm import declarative_base
+from pg_rls import rls, policy, Policy, PolicyType, PolicyCommands
+
+BaseModel = declarative_base()
+
+@rls()
+@policy(Policy("pol_my_models_select_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.SELECT, using="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_delete_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.DELETE, using="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_update_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.UPDATE, using="user_id == auth.uid()", with_check="user_id == auth.uid()"))
+@policy(Policy("pol_my_models_update_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.INSERT, with_check="user_id == auth.uid()"))
+# Equivalent to:
+# @policy(Policy("pol_my_models_primary", as_=PolicyType.PERMISSIVE, for_=PolicyCommands.ALL, using="user_id == auth.uid()", with_check="user_id == auth.uid()"))
+class MyModel(BaseModel):
+    ...
+```
+

pgalchemy-0.1.2/pgalchemy/__init__.py

@@ -0,0 +1,2 @@
+from .sqlalchemy import rls_base, rls, policy
+from .policy import PolicyType, PolicyCommands, Policy

pgalchemy-0.1.2/pgalchemy/alembic/__init__.py

@@ -0,0 +1 @@
+from .operations import EnableRlsOp, DisableRlsOp, DropPolicyOp, CreatePolicyOp, AlterPolicyOp, RenamePolicyOp, PolicySql

pgalchemy-0.1.2/pgalchemy/alembic/comparator.py

@@ -0,0 +1,39 @@
+from alembic.autogenerate import comparators
+from alembic.operations import Operations, MigrateOperation
+from sqlalchemy import Table
+
+from pg_rls.sqlalchemy import RlsData
+from .operations import EnableRlsOp, DisableRlsOp
+from .. import Policy
+
+
+@comparators.dispatch_for("table")
+def compare_rls(autogen_context, modify_ops, schemaname, tablename, conn_table, metadata_table: Table):
+    rls: RlsData = metadata_table.info.get('rls')
+
+    db_table, rls_enabled_db = get_table_rls_data(autogen_context, schemaname, tablename)
+    if db_table is None:
+        return
+
+    compare_rls_enabled(modify_ops, rls, rls_enabled_db, schemaname, tablename)
+
+
+def compare_rls_enabled(modify_ops, rls, rls_enabled_db, schemaname, tablename):
+    if rls.active is True and rls_enabled_db is False:
+        modify_ops.ops.append(
+            EnableRlsOp(tablename, schema=schemaname)
+        )
+    if rls.active is False and rls_enabled_db is True:
+        modify_ops.ops.append(
+            DisableRlsOp(tablename, schema=schemaname)
+        )
+
+
+def get_table_rls_data(autogen_context, schemaname, tablename):
+    results = autogen_context.connection.execute(
+        'select relname, relrowsecurity, relforcerowsecurity from pg_class where  relnamespace = %s and relname = %s;',
+        (schemaname, tablename)
+    )
+    db_table = results.fetchone()
+    rls_enabled_db = db_table['relrowsecurity'] if db_table is not None else None
+    return db_table, rls_enabled_db

pgalchemy-0.1.2/pgalchemy/alembic/operations/__init__.py

@@ -0,0 +1 @@
+from .rls import DisableRlsOp, EnableRlsOp

pgalchemy-0.1.2/pgalchemy/alembic/operations/rls.py

@@ -0,0 +1,57 @@
+from alembic.operations import MigrateOperation, Operations
+
+
+@Operations.register_operation("enable_rls")
+class EnableRlsOp(MigrateOperation):
+    """Enable RLS on a table."""
+
+    def __init__(self, table_name, schema=None):
+        self.table_name = table_name
+        self.schema = schema
+
+    @classmethod
+    def enable_rls(cls, operations, table_name, **kw):
+        """Issue a "ALTER TABLE ENABLE ROW LEVEL SECURITY" instruction."""
+
+        op = EnableRlsOp(table_name, **kw)
+        return operations.invoke(op)
+
+    def reverse(self):
+        # only needed to support autogenerate
+        return DisableRlsOp(self.table_name, schema=self.schema)
+
+
+@Operations.register_operation("disable_rls")
+class DisableRlsOp(MigrateOperation):
+    """Disable RLS on table."""
+
+    def __init__(self, table_name, schema=None):
+        self.table_name = table_name
+        self.schema = schema
+
+    @classmethod
+    def disable_rls(cls, operations, sequence_name, **kw):
+        """Issue a "ALTER TABLE DISABLE ROW LEVEL SECURITY" instruction."""
+
+        op = DisableRlsOp(sequence_name, **kw)
+        return operations.invoke(op)
+
+    def reverse(self):
+        # only needed to support autogenerate
+        return EnableRlsOp(self.sequence_name, schema=self.schema)
+
+@Operations.implementation_for(EnableRlsOp)
+def enable_rls(operations, operation: EnableRlsOp):
+    if operation.schema is not None:
+        name = "%s.%s" % (operation.schema, operation.table_name)
+    else:
+        name = operation.table_name
+    operations.execute("ALTER TABLE %s ENABLE ROW LEVEL SECURITY;" % name)
+
+@Operations.implementation_for(DisableRlsOp)
+def disable_rls(operations, operation: DisableRlsOp):
+    if operation.schema is not None:
+        name = "%s.%s" % (operation.schema, operation.table_name)
+    else:
+        name = operation.table_name
+    operations.execute("ALTER TABLE %s DISABLE ROW LEVEL SECURITY;" % name)

pgalchemy-0.1.2/pgalchemy/policy.py

@@ -0,0 +1,43 @@
+from enum import Enum
+
+from alembic_utils.pg_policy import PGPolicy
+from sqlalchemy import Table, BinaryExpression
+
+
+class PolicyType(Enum):
+    PERMISSIVE = 'PERMISSIVE'
+    RESTRICTIVE = 'RESTRICTIVE'
+
+
+class PolicyCommands(Enum):
+    ALL = 'ALL'
+    SELECT = 'SELECT'
+    INSERT = 'INSERT'
+    UPDATE = 'UPDATE'
+    DELETE = 'DELETE'
+
+
+class Policy:
+    def __init__(
+            self,
+            name: str,
+            as_: PolicyType = PolicyType.PERMISSIVE,
+            for_: PolicyCommands = PolicyCommands.ALL,
+            using: str | BinaryExpression | None = None,
+            with_check: str | BinaryExpression | None = None
+    ):
+        self.name = name
+        self.with_check = str(with_check)
+        self.using = str(using)
+        self.for_ = for_
+        self.as_ = as_
+
+    @classmethod
+    def from_pg_policy(cls, row):
+        return Policy(
+            name=row['policyname'],
+            as_=PolicyType(row['permissive']),
+            for_=PolicyCommands(row['cmd']),
+            using=row['qual'],
+            with_check=row['with_check'],
+        )

pgalchemy-0.1.2/pgalchemy/policy_sql.py

@@ -0,0 +1,56 @@
+from sqlalchemy import Table
+
+from .policy import Policy
+
+
+class PolicySql:
+    def __init__(self, policy: Policy, table_name: str, schema: str):
+        self.policy = policy
+        self.table_name = table_name
+        self.schema = schema
+    
+    def qualify(self, name):
+        return f"{self.schema}.{name}"
+    
+    @property
+    def tab(self):
+        return self.qualify(self.table_name)
+    
+    @property
+    def pol(self):
+        return self.qualify(self.policy.name)
+
+    # Full
+    def rename_sql(self, old_name):
+        return f"alter policy {old_name} on {self.tab} rename to {self.pol}"
+
+    def drop_sql(self):
+        return f"drop policy {self.pol} on {self.tab}"
+
+    def create_sql(self):
+        return self._create_fragment() + self._as_fragment() + self._from_fragment() + self._using_fragment() + self._with_check_fragment()
+
+    def alter_sql(self):
+        return self._alter_fragment() + self._using_fragment() + self._with_check_fragment()
+
+    # Fragments
+    def _create_fragment(self):
+        return f"create policy {self.pol} on {self.tab}\n"
+
+    def _alter_fragment(self):
+        return f"alter policy {self.pol} on {self.tab}\n"
+
+    def _as_fragment(self):
+        return f"as {self.policy.as_}\n" if self.policy.as_ else ""
+
+    def _from_fragment(self):
+        return f"for {self.policy.for_}\n" if self.policy.for_ else ""
+
+    def _with_check_fragment(self):
+        return f"with check {self.policy.with_check}\n" if self.policy.with_check else ""
+
+    def _using_fragment(self):
+        return f"using {self.policy.using}\n" if self.policy.using else ""
+
+    def definition_sql(self):
+        return self._as_fragment() + self._from_fragment() + self._using_fragment() + self._with_check_fragment()

pgalchemy-0.1.2/pgalchemy/sqlalchemy.py

@@ -0,0 +1,66 @@
+from typing import Type, List, Optional
+
+from alembic_utils.pg_policy import PGPolicy
+from sqlalchemy import Table
+from sqlalchemy.event import listens_for
+from sqlalchemy.orm import DeclarativeBase, Mapper
+
+from .policy import Policy, PolicyType
+from .policy_sql import PolicySql
+
+
+class RlsData:
+    def __init__(self, active):
+        self.active = active
+        self.policies: List[Policy] = []
+
+
+def rls_base(Base: Type[DeclarativeBase], default_active: bool = True):
+    class WithRls(Base):
+        __rls__ = RlsData(default_active)
+
+    @listens_for(WithRls, 'after_configured')
+    def receive_mapper_configured(mapper: Mapper, class_: Type[WithRls]):
+        table: Table = mapper.mapped_table()
+        rls = getattr(class_, '__rls__')
+        table.info.setdefault('rls', rls)
+        if rls.active:
+            for policy in rls.policies:
+                attach_policy(policy, table)
+
+    return WithRls
+
+
+def rls(enabled=True, policies: Optional[List[Policy]] = None):
+    def wrapper(Model: Type[DeclarativeBase]):
+        Model.__rls__ = RlsData(enabled)
+        Model.__rls__.policies = policies or []
+        return Model
+    return wrapper
+
+
+def rls_for_table(enabled=True, policies: Optional[List[Policy]] = None):
+    def wrapper(table: Table):
+        data = RlsData(enabled)
+        data.policies = policies or []
+        table.info.setdefault('rls', data)
+        return Table
+    return wrapper
+
+
+def policy(pol: Policy):
+    def wrapper(Model: Type[DeclarativeBase]):
+        if not Model.__rls__:
+            Model.__rls__ = RlsData(True)
+        Model.__rls__.policies.append(pol)
+        return Model
+    return wrapper
+
+
+def attach_policy(policy: Policy, table: Table):
+    return PGPolicy(
+        on_entity=table.name,
+        schema=table.schema,
+        signature=policy.name,
+        definition=PolicySql(policy, table.name, table.schema)
+    )

pgalchemy-0.1.2/pyproject.toml

@@ -0,0 +1,21 @@
+[tool.poetry]
+name = "pgalchemy"
+version = "0.1.2"
+description = ""
+authors = ["Matthew Beatty <beattyml1@gmail.com>"]
+readme = "README.md"
+include = [{path="README.md"}, {path="LICENSE"}]
+packages = [
+    { include = "pgalchemy", from = "." },
+]
+
+[tool.poetry.dependencies]
+python = "^3.11"
+sqlalchemy = "^2.0.32"
+alembic = "^1.13.2"
+alembic-utils = "^0.8.4"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"