-- ============================================================================
-- Free-form SQL Collaboration Demo: Data Provider
-- ============================================================================
-- This example demonstrates a Snowflake Data Clean Rooms collaboration using
-- freeform SQL policies. The data provider creates a sample dataset with
-- Snowflake aggregation, join, and masking policies, registers a data offering
-- that permits freeform SQL queries, creates a template, and initializes a
-- collaboration with one other collaborator (data_consumer).
--
-- For more information, see:
--   docs.snowflake.com/user-guide/cleanrooms/free-form-sql.rst
--   docs.snowflake.com/user-guide/cleanrooms/spec-reference
-- ============================================================================


-- ============================================================================
-- SETUP: Create sample database, schema, table, and policies.
-- ============================================================================

USE ROLE SAMOOHA_APP_ROLE;
USE WAREHOUSE APP_WH;

-- You can't use secondary roles with most collaboration procedures.
USE SECONDARY ROLES NONE;

CREATE DATABASE IF NOT EXISTS PROVIDER_DB;
CREATE SCHEMA IF NOT EXISTS PROVIDER_DB.DATA_SCH;

-- Create a table with 300 rows from the sample CUSTOMERS table.
CREATE OR REPLACE TABLE PROVIDER_DB.DATA_SCH.CUSTOMERS AS
  SELECT HASHED_EMAIL, STATUS, AGE_BAND, REGION_CODE, DAYS_ACTIVE, INCOME_BRACKET
  FROM SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS
  LIMIT 300;

-- Create an aggregation policy that requires a minimum group size of 5.
CREATE OR REPLACE AGGREGATION POLICY PROVIDER_DB.DATA_SCH.MIN_GROUP_SIZE_POLICY
  AS () RETURNS AGGREGATION_CONSTRAINT ->
    AGGREGATION_CONSTRAINT(MIN_GROUP_SIZE => 5);

-- Create an inactive join policy. You will modify this later.
CREATE OR REPLACE JOIN POLICY PROVIDER_DB.DATA_SCH.EMAIL_JOIN_POLICY
  AS () RETURNS JOIN_CONSTRAINT ->
    JOIN_CONSTRAINT(JOIN_REQUIRED => FALSE);

-- Create a masking policy that replaces the original value with a fixed string.
CREATE OR REPLACE MASKING POLICY PROVIDER_DB.DATA_SCH.MASK_INCOME_POLICY
  AS (val STRING) RETURNS STRING ->
    '***MASKED***';


-- ============================================================================
-- Register a data offering with freeform SQL policies.
-- ============================================================================

-- The data offering enables freeform SQL queries (template_and_freeform_sql)
-- and attaches three Snowflake policies to protect data in freeform queries:
--   * Aggregation policy on hashed_email: enforces a minimum group size of 5.
--   * Join policy on hashed_email: requires joins to include this column.
--   * Masking policy on income_bracket: masks the column value in query results.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_DATA_OFFERING(
  $$
  api_version: 2.0.0
  spec_type: data_offering
  version: V1
  name: provider_customers
  description: Customer dataset with freeform SQL policies.
  datasets:
    - alias: customers
      data_object_fqn: PROVIDER_DB.DATA_SCH.CUSTOMERS
      object_class: custom
      allowed_analyses: template_and_freeform_sql
      schema_and_template_policies:
        HASHED_EMAIL:
          category: join_standard
          column_type: hashed_email_b64_encoded
        STATUS:
          category: passthrough
        AGE_BAND:
          category: passthrough
        DAYS_ACTIVE:
          category: passthrough
        INCOME_BRACKET:
          category: passthrough
      freeform_sql_policies:
        aggregation_policy:
          name: PROVIDER_DB.DATA_SCH.MIN_GROUP_SIZE_POLICY
          entity_keys:
            - HASHED_EMAIL_B64_ENCODED
        join_policy:
          name: PROVIDER_DB.DATA_SCH.EMAIL_JOIN_POLICY
          columns:
            - HASHED_EMAIL_B64_ENCODED
        masking_policies:
          - name: PROVIDER_DB.DATA_SCH.MASK_INCOME_POLICY
            columns:
              - INCOME_BRACKET
  $$
);

-- Save the data offering ID returned by the registration call.
SET data_offering_id = '<data_offering_id>';

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.VIEW_REGISTERED_DATA_OFFERINGS();


-- ============================================================================
-- Register a template with a simple one-table query.
-- ============================================================================

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
  $$
  api_version: 2.0.0
  spec_type: template
  name: status_summary
  version: V1
  type: sql_analysis
  description: Returns a count of customers grouped by status.
  template:
    SELECT status, COUNT(*) AS customer_count
      FROM IDENTIFIER({{ source_table[0] }})
      GROUP BY status
      ORDER BY customer_count DESC;
  $$
);

-- Save the template ID returned by the registration call.
SET template_id = '<template_id>';

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.VIEW_REGISTERED_TEMPLATES();


-- ============================================================================
-- Create the collaboration.
-- ============================================================================

-- Replace the <...> placeholders with the appropriate values.
-- Get your account data sharing ID with:
--   SELECT CURRENT_ORGANIZATION_NAME() || '.' || CURRENT_ACCOUNT_NAME();
-- In this collaboration, the consumer can run templated and free-form queries
-- against the provider's data. The provider/owner isn't an analysis runner.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.INITIALIZE(
  $$
  api_version: 2.0.0
  spec_type: collaboration
  name: freeform_sql_demo
  owner: data_provider
  collaborator_identifier_aliases:
    data_provider: <provider_account_data_sharing_id>
    data_consumer: <consumer_account_data_sharing_id>
  analysis_runners:
    data_consumer:
      data_providers:
        data_provider:
          data_offerings:
            - id: <data_offering_id>
      templates:
        - id: <template_id>
  $$,
  'APP_WH'
);

SET collaboration_name = 'freeform_sql_demo';

-- INITIALIZE automatically joins the owner. Repeat until status is JOINED.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.GET_STATUS($collaboration_name);

-- Verify that the collaboration is visible.
-- Collaboration spec is in COLLABORATION_SPEC column.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.VIEW_COLLABORATIONS() ->>
  SELECT * FROM $1 WHERE "SOURCE_NAME" = $collaboration_name;

-- SWITCH TO data_consumer account to join and run analyses.

-- Update the join policy associated with HASHED_EMAIL_B64_ENCODED.
-- All queries on that data offering now require joins on HASHED_EMAIL_B64_ENCODED.
-- Re-run any of the previously successful free-form queries and they will fail.
ALTER JOIN POLICY PROVIDER_DB.DATA_SCH.EMAIL_JOIN_POLICY SET BODY ->
  JOIN_CONSTRAINT(JOIN_REQUIRED => TRUE);

-- ============================================================================
-- CLEANUP: Delete the collaboration, registered resources, and sample data.
-- ============================================================================

-- Teardown is a multi-step process. Call TEARDOWN, then wait for GET_STATUS
-- to report LOCAL_DROP_PENDING, then call TEARDOWN again.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.TEARDOWN($collaboration_name);
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.GET_STATUS($collaboration_name);

-- When GET_STATUS reports LOCAL_DROP_PENDING, call TEARDOWN again to complete.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.TEARDOWN($collaboration_name);
