Snowflake Data Lake + Azure Integration with Programmatic Access Token (PAT) Authentication

-- ============================================
-- SET THESE VARIABLES BEFORE RUNNING
-- ============================================
SET db_name = 'OPENBRIDGE';
SET schema_name = 'OPENBRIDGE_SCHEMA';
SET user_name = 'OPENBRIDGE_USER';
SET role_name = 'OPENBRIDGE_ROLE';
SET stage_name = 'OPENBRIDGE_STAGE';
SET warehouse_name = 'OPENBRIDGE_WH';
SET warehouse_type = 'STANDARD';
SET warehouse_size = 'XSMALL';
SET azure_storage_account = 'your-storage-account';
SET azure_container_name = 'your-container-name';
SET azure_tenant_id = 'your-tenant-id';

SET default_ns = $db_name || '.' || $schema_name;
SET azure_url = 'azure://' || $azure_storage_account || '.blob.core.windows.net/' || $azure_container_name || '/';

-- ============================================
-- CREATE ROLE
-- ============================================
USE ROLE ACCOUNTADMIN;
CREATE ROLE IF NOT EXISTS IDENTIFIER($role_name);

-- Grant the role to your current user (for testing purposes)
SET current_username = CURRENT_USER();
GRANT ROLE IDENTIFIER($role_name) TO USER IDENTIFIER($current_username);

-- ============================================
-- CREATE WAREHOUSE
-- ============================================
CREATE WAREHOUSE IF NOT EXISTS IDENTIFIER($warehouse_name)
  WAREHOUSE_SIZE = $warehouse_size
  WAREHOUSE_TYPE = $warehouse_type
  AUTO_SUSPEND = 60
  AUTO_RESUME = TRUE
  INITIALLY_SUSPENDED = TRUE;

GRANT USAGE ON WAREHOUSE IDENTIFIER($warehouse_name) TO ROLE IDENTIFIER($role_name);

-- ============================================
-- CREATE DATABASE AND SCHEMA
-- ============================================
CREATE DATABASE IF NOT EXISTS IDENTIFIER($db_name);
GRANT USAGE ON DATABASE IDENTIFIER($db_name) TO ROLE IDENTIFIER($role_name);

USE DATABASE IDENTIFIER($db_name);
CREATE SCHEMA IF NOT EXISTS IDENTIFIER($schema_name);

GRANT ALL PRIVILEGES ON SCHEMA IDENTIFIER($schema_name) TO ROLE IDENTIFIER($role_name);
GRANT CREATE TABLE, CREATE VIEW, CREATE STAGE, CREATE FILE FORMAT ON SCHEMA IDENTIFIER($schema_name) TO ROLE IDENTIFIER($role_name);
USE SCHEMA IDENTIFIER($default_ns);

-- ============================================
-- CREATE USER
-- Service users cannot have passwords — authentication
-- is handled exclusively via the PAT created in Step 4.
-- ============================================
CREATE USER IF NOT EXISTS IDENTIFIER($user_name)
  TYPE = SERVICE
  DEFAULT_ROLE = $role_name
  DEFAULT_WAREHOUSE = $warehouse_name
  DEFAULT_NAMESPACE = $default_ns;

GRANT ROLE IDENTIFIER($role_name) TO USER IDENTIFIER($user_name);

-- ============================================
-- CREATE AZURE STORAGE INTEGRATION AND STAGE
-- (AZURE_TENANT_ID, STORAGE_ALLOWED_LOCATIONS, and URL
-- do not support session variable substitution)
-- ============================================
DECLARE
  tenant_id VARCHAR DEFAULT $azure_tenant_id;
  azure_uri VARCHAR DEFAULT $azure_url;
  stage VARCHAR DEFAULT $stage_name;
  sql_str VARCHAR;
BEGIN
  sql_str := 'CREATE OR REPLACE STORAGE INTEGRATION openbridge_azure TYPE = EXTERNAL_STAGE STORAGE_PROVIDER = ''AZURE'' ENABLED = TRUE AZURE_TENANT_ID = ''' || :tenant_id || ''' STORAGE_ALLOWED_LOCATIONS = (''' || :azure_uri || ''')';
  EXECUTE IMMEDIATE :sql_str;
  sql_str := 'CREATE STAGE IF NOT EXISTS ' || :stage || ' STORAGE_INTEGRATION = openbridge_azure URL = ''' || :azure_uri || ''' FILE_FORMAT = (TYPE = ''PARQUET'')';
  EXECUTE IMMEDIATE :sql_str;
END;

GRANT ALL PRIVILEGES ON STAGE IDENTIFIER($stage_name) TO ROLE IDENTIFIER($role_name);
GRANT USAGE ON INTEGRATION openbridge_azure TO ROLE IDENTIFIER($role_name);

-- ============================================
-- CREATE NETWORK POLICY
-- ============================================
USE ROLE ACCOUNTADMIN;

-- These are the Openbridge production IP addresses.
-- Current IPs: https://docs.openbridge.com/en/articles/11822520-snowflake-warehouse-with-oauth-authentication
-- If these IPs ever change, update the network rule:
--   ALTER NETWORK RULE openbridge_network_rule
--     SET VALUE_LIST = ('NEW_IP_1/32', 'NEW_IP_2/32');
CREATE NETWORK RULE IF NOT EXISTS openbridge_network_rule
  TYPE = IPV4
  MODE = INGRESS
  VALUE_LIST = ('52.2.68.68/32', '52.54.227.22/32');

CREATE NETWORK POLICY IF NOT EXISTS openbridge_network_policy
  ALLOWED_NETWORK_RULE_LIST = ('openbridge_network_rule');

ALTER USER IDENTIFIER($user_name)
  SET NETWORK_POLICY = openbridge_network_policy;

-- ============================================
-- CREATE AUTHENTICATION POLICY FOR PAT
-- ============================================
CREATE AUTHENTICATION POLICY IF NOT EXISTS openbridge_pat_policy
  PAT_POLICY = (
    NETWORK_POLICY_EVALUATION = ENFORCED_REQUIRED,
    MAX_EXPIRY_IN_DAYS = 365
  );

ALTER USER IDENTIFIER($user_name)
  SET AUTHENTICATION POLICY openbridge_pat_policy;

-- ============================================
-- GRANT PAT MANAGEMENT TO OPENBRIDGE ROLE
-- ============================================
GRANT MODIFY PROGRAMMATIC AUTHENTICATION METHODS ON USER IDENTIFIER($user_name)
  TO ROLE IDENTIFIER($role_name);

-- Test 1: Create a table
CREATE OR REPLACE TABLE test_table (
  id INT AUTOINCREMENT PRIMARY KEY,
  name STRING,
  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP()
);

-- Test 2: Insert data into the table
INSERT INTO test_table (name) VALUES
  ('Alice'),
  ('Bob'),
  ('Charlie');

-- Test 3: Query data from the table
SELECT * FROM test_table;

-- Test 4: Update data in the table
UPDATE test_table
SET name = 'Alice Updated'
WHERE name = 'Alice';

-- Test 5: Delete data from the table
DELETE FROM test_table
WHERE name = 'Bob';

-- Test 6: Drop the table
DROP TABLE test_table;

GRANT OWNERSHIP ON DATABASE IDENTIFIER($db_name) TO ROLE IDENTIFIER($role_name);
GRANT OWNERSHIP ON SCHEMA IDENTIFIER($schema_name) TO ROLE IDENTIFIER($role_name);

DESC STORAGE INTEGRATION openbridge_azure;

SE ROLE <YOUR_ROLE_NAME>;

ALTER USER <YOUR_USER_NAME>
  ADD PROGRAMMATIC ACCESS TOKEN openbridge_pat
    DAYS_TO_EXPIRY = 365
    ROLE_RESTRICTION = <YOUR_ROLE_NAME>;

ALTER USER OPENBRIDGE_USER
  ADD PROGRAMMATIC ACCESS TOKEN openbridge_pat
    DAYS_TO_EXPIRY = 365
    ROLE_RESTRICTION = OPENBRIDGE_ROLE;

SELECT CURRENT_ORGANIZATION_NAME() || '-' || CURRENT_ACCOUNT_NAME();

SHOW PROGRAMMATIC ACCESS TOKENS FOR USER <YOUR_USER_NAME>;

ALTER USER <YOUR_USER_NAME>
  REMOVE PROGRAMMATIC ACCESS TOKEN openbridge_pat;

USE ROLE <YOUR_ROLE_NAME>;

ALTER USER <YOUR_USER_NAME>
  ADD PROGRAMMATIC ACCESS TOKEN openbridge_pat
    DAYS_TO_EXPIRY = 365
    ROLE_RESTRICTION = <YOUR_ROLE_NAME>;

ALTER AUTHENTICATION POLICY your_existing_policy
  SET AUTHENTICATION_METHODS = ('OAUTH', 'PASSWORD', 'PROGRAMMATIC_ACCESS_TOKEN');

DESCRIBE USER OPENBRIDGE_USER;