-- ============================================================================
-- SQL Server Row-Level Security (RLS) Setup
-- ============================================================================
-- This script sets up database-level RLS as the last line of defense
-- Use with ORM middleware for defense-in-depth security

USE ORMtest;
GO

-- ============================================================================
-- 1. Add tenant_id column to tables (if not exists)
-- ============================================================================

-- Check and add tenant_id to Customers
IF NOT EXISTS (SELECT * FROM sys.columns WHERE object_id = OBJECT_ID('dbo.Customers') AND name = 'tenant_id')
BEGIN
    ALTER TABLE dbo.Customers ADD tenant_id INT NOT NULL DEFAULT 1;
END
GO

-- Check and add tenant_id to Orders
IF NOT EXISTS (SELECT * FROM sys.columns WHERE object_id = OBJECT_ID('dbo.Orders') AND name = 'tenant_id')
BEGIN
    ALTER TABLE dbo.Orders ADD tenant_id INT NOT NULL DEFAULT 1;
END
GO

-- Check and add tenant_id to Products
IF NOT EXISTS (SELECT * FROM sys.columns WHERE object_id = OBJECT_ID('dbo.Products') AND name = 'tenant_id')
BEGIN
    ALTER TABLE dbo.Products ADD tenant_id INT NOT NULL DEFAULT 1;
END
GO

-- ============================================================================
-- 2. Create Security Predicate Functions
-- ============================================================================

-- Function for tenant isolation
CREATE OR ALTER FUNCTION dbo.fn_tenantSecurityPredicate(@TenantId INT)
    RETURNS TABLE
    WITH SCHEMABINDING
AS
    RETURN SELECT 1 AS fn_securitypredicate_result
    WHERE @TenantId = CAST(SESSION_CONTEXT(N'TenantId') AS INT)
        OR SESSION_CONTEXT(N'TenantId') IS NULL; -- Allow queries when context not set
GO

-- Function for user ownership (Orders)
CREATE OR ALTER FUNCTION dbo.fn_userOwnershipPredicate(@UserId INT, @TenantId INT)
    RETURNS TABLE
    WITH SCHEMABINDING
AS
    RETURN SELECT 1 AS fn_securitypredicate_result
    WHERE (
        @UserId = CAST(SESSION_CONTEXT(N'UserId') AS INT)
        AND @TenantId = CAST(SESSION_CONTEXT(N'TenantId') AS INT)
    )
    OR CAST(SESSION_CONTEXT(N'Role') AS NVARCHAR(50)) = 'admin'
    OR SESSION_CONTEXT(N'UserId') IS NULL; -- Allow queries when context not set
GO

-- Function for role-based access (Payments)
CREATE OR ALTER FUNCTION dbo.fn_roleBasedPredicate(@TenantId INT)
    RETURNS TABLE
    WITH SCHEMABINDING
AS
    RETURN SELECT 1 AS fn_securitypredicate_result
    WHERE @TenantId = CAST(SESSION_CONTEXT(N'TenantId') AS INT)
        AND CAST(SESSION_CONTEXT(N'Role') AS NVARCHAR(50)) IN ('admin', 'finance')
        OR SESSION_CONTEXT(N'TenantId') IS NULL; -- Allow queries when context not set
GO

-- ============================================================================
-- 3. Create Security Policies
-- ============================================================================

-- Drop existing policies if they exist
DROP SECURITY POLICY IF EXISTS dbo.CustomersTenantPolicy;
DROP SECURITY POLICY IF EXISTS dbo.OrdersTenantPolicy;
DROP SECURITY POLICY IF EXISTS dbo.ProductsTenantPolicy;
GO

-- Customers table - tenant isolation
CREATE SECURITY POLICY dbo.CustomersTenantPolicy
    ADD FILTER PREDICATE dbo.fn_tenantSecurityPredicate(tenant_id)
        ON dbo.Customers,
    ADD BLOCK PREDICATE dbo.fn_tenantSecurityPredicate(tenant_id)
        ON dbo.Customers AFTER INSERT,
    ADD BLOCK PREDICATE dbo.fn_tenantSecurityPredicate(tenant_id)
        ON dbo.Customers AFTER UPDATE
    WITH (STATE = ON);
GO

-- Orders table - user ownership
CREATE SECURITY POLICY dbo.OrdersTenantPolicy
    ADD FILTER PREDICATE dbo.fn_userOwnershipPredicate(CustomerId, tenant_id)
        ON dbo.Orders,
    ADD BLOCK PREDICATE dbo.fn_userOwnershipPredicate(CustomerId, tenant_id)
        ON dbo.Orders AFTER INSERT,
    ADD BLOCK PREDICATE dbo.fn_userOwnershipPredicate(CustomerId, tenant_id)
        ON dbo.Orders AFTER UPDATE
    WITH (STATE = ON);
GO

-- Products table - tenant isolation
CREATE SECURITY POLICY dbo.ProductsTenantPolicy
    ADD FILTER PREDICATE dbo.fn_tenantSecurityPredicate(tenant_id)
        ON dbo.Products,
    ADD BLOCK PREDICATE dbo.fn_tenantSecurityPredicate(tenant_id)
        ON dbo.Products AFTER INSERT,
    ADD BLOCK PREDICATE dbo.fn_tenantSecurityPredicate(tenant_id)
        ON dbo.Products AFTER UPDATE
    WITH (STATE = ON);
GO

-- ============================================================================
-- 4. Test RLS Setup
-- ============================================================================

-- Set context for tenant 1
EXEC sp_set_session_context 'TenantId', 1;
EXEC sp_set_session_context 'UserId', 123;
EXEC sp_set_session_context 'Role', 'user';
GO

-- This query will only return records where tenant_id = 1
SELECT TOP 10 * FROM dbo.Customers;
GO

-- Reset context
EXEC sp_set_session_context 'TenantId', NULL;
EXEC sp_set_session_context 'UserId', NULL;
EXEC sp_set_session_context 'Role', NULL;
GO

-- ============================================================================
-- 5. Integration with ORM
-- ============================================================================

/*
In your application code, set context before each query:

SQL Server adapter example:

// Before running queries
await connection.request()
    .query(`EXEC sp_set_session_context 'TenantId', ${user.tenantId}`);
await connection.request()
    .query(`EXEC sp_set_session_context 'UserId', ${user.id}`);
await connection.request()
    .query(`EXEC sp_set_session_context 'Role', '${user.role}'`);

// Then run your ORM queries - RLS is enforced automatically
const customers = await client.Customers.findMany({});

// After queries (optional cleanup)
await connection.request()
    .query(`EXEC sp_set_session_context 'TenantId', NULL`);
*/

-- ============================================================================
-- 6. Verify Policies are Active
-- ============================================================================

-- List all security policies
SELECT 
    SCHEMA_NAME(schema_id) AS schema_name,
    name AS policy_name,
    is_enabled,
    is_schema_bound
FROM sys.security_policies;
GO

-- List all security predicates
SELECT 
    OBJECT_NAME(sp.object_id) AS policy_name,
    OBJECT_NAME(spp.target_object_id) AS table_name,
    spp.predicate_type_desc,
    OBJECT_NAME(spp.predicate_definition_function_id) AS predicate_function
FROM sys.security_policies sp
INNER JOIN sys.security_predicates spp
    ON sp.object_id = spp.object_id;
GO

-- ============================================================================
-- 7. Drop Policies and Functions (for cleanup/testing)
-- ============================================================================

/*
-- Uncomment to drop policies:
DROP SECURITY POLICY IF EXISTS dbo.CustomersTenantPolicy;
DROP SECURITY POLICY IF EXISTS dbo.OrdersTenantPolicy;
DROP SECURITY POLICY IF EXISTS dbo.ProductsTenantPolicy;
GO

-- Uncomment to drop functions:
DROP FUNCTION IF EXISTS dbo.fn_tenantSecurityPredicate;
DROP FUNCTION IF EXISTS dbo.fn_userOwnershipPredicate;
DROP FUNCTION IF EXISTS dbo.fn_roleBasedPredicate;
GO

-- Uncomment to remove tenant_id columns:
ALTER TABLE dbo.Customers DROP COLUMN IF EXISTS tenant_id;
ALTER TABLE dbo.Orders DROP COLUMN IF EXISTS tenant_id;
ALTER TABLE dbo.Products DROP COLUMN IF EXISTS tenant_id;
GO
*/
