#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

import codecs
import os
import time
import uuid
from urllib.parse import unquote

import boto3
import botocore
import pytest
from py4j.protocol import Py4JJavaError


from iceberg_spark import IcebergSparkSession
from apache_polaris.sdk.catalog import CreateNamespaceRequest, CreateTableRequest, ModelSchema
from apache_polaris.sdk.catalog.api.iceberg_catalog_api import IcebergCatalogAPI
from apache_polaris.sdk.catalog.api.iceberg_o_auth2_api import IcebergOAuth2API
from apache_polaris.sdk.catalog.api_client import ApiClient as CatalogApiClient
from apache_polaris.sdk.catalog.configuration import Configuration
from apache_polaris.sdk.management import ApiClient as ManagementApiClient
from apache_polaris.sdk.management import PolarisDefaultApi, Principal, PrincipalRole, CatalogRole, \
  CatalogGrant, CatalogPrivilege, ApiException, CreateCatalogRoleRequest, CreatePrincipalRoleRequest, \
  CreatePrincipalRequest, AddGrantRequest, GrantCatalogRoleRequest, GrantPrincipalRoleRequest, UpdateCatalogRequest


@pytest.fixture
def snowman(polaris_url, polaris_catalog_url, root_client, snowflake_catalog):
  """
  create the snowman principal with full table/namespace privileges
  :param root_client:
  :param snowflake_catalog:
  :return:
  """
  snowman_name = f"snowman_{str(uuid.uuid4())[-10:]}"
  table_writer_rolename = f"table_writer_{str(uuid.uuid4())[-10:]}"
  snowflake_writer_rolename = "snowflake_writer"
  try:
    snowman = create_principal(polaris_url, polaris_catalog_url, root_client, snowman_name)
    writer_principal_role = create_principal_role(root_client, table_writer_rolename)
    writer_catalog_role = create_catalog_role(root_client, snowflake_catalog, snowflake_writer_rolename)
    root_client.assign_catalog_role_to_principal_role(principal_role_name=writer_principal_role.name,
                                                      catalog_name=snowflake_catalog.name,
                                                      grant_catalog_role_request=GrantCatalogRoleRequest(
                                                        catalog_role=writer_catalog_role))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, writer_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.TABLE_FULL_METADATA)))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, writer_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.VIEW_FULL_METADATA)))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, writer_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.TABLE_WRITE_DATA)))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, writer_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.NAMESPACE_FULL_METADATA)))

    root_client.assign_principal_role(snowman.principal.name,
                                      grant_principal_role_request=GrantPrincipalRoleRequest(
                                        principal_role=writer_principal_role))
    yield snowman
  finally:
    root_client.delete_principal(snowman_name)
    root_client.delete_principal_role(principal_role_name=table_writer_rolename)
    root_client.delete_catalog_role(catalog_role_name=snowflake_writer_rolename, catalog_name=snowflake_catalog.name)


@pytest.fixture
def reader(polaris_url, polaris_catalog_url, root_client, snowflake_catalog):
  """
  create the test_reader principal with table/namespace list and read privileges

  :param root_client:
  :param snowflake_catalog:
  :return:
  """
  reader_principal_name = 'test_reader'
  reader_principal_role_name = "table_reader"
  reader_catalog_role_name = 'snowflake_reader'
  try:
    reader = create_principal(polaris_url, polaris_catalog_url, root_client, reader_principal_name)
    reader_principal_role = create_principal_role(root_client, reader_principal_role_name)
    reader_catalog_role = create_catalog_role(root_client, snowflake_catalog, reader_catalog_role_name)

    root_client.assign_catalog_role_to_principal_role(principal_role_name=reader_principal_role.name,
                                                      catalog_name=snowflake_catalog.name,
                                                      grant_catalog_role_request=GrantCatalogRoleRequest(
                                                        catalog_role=reader_catalog_role))
    root_client.assign_principal_role(reader.principal.name,
                                      grant_principal_role_request=GrantPrincipalRoleRequest(
                                        principal_role=reader_principal_role))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, reader_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.TABLE_READ_DATA)))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, reader_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.TABLE_LIST)))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, reader_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.TABLE_READ_PROPERTIES)))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, reader_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.NAMESPACE_LIST)))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, reader_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.NAMESPACE_READ_PROPERTIES)))
    yield reader
  finally:
    root_client.delete_principal(reader_principal_name)
    root_client.delete_principal_role(principal_role_name=reader_principal_role_name)
    root_client.delete_catalog_role(catalog_role_name=reader_catalog_role_name, catalog_name=snowflake_catalog.name)


@pytest.fixture
def snowman_catalog_client(polaris_catalog_url, snowman):
  """
  Create an iceberg catalog client with snowman credentials
  :param polaris_catalog_url:
  :param snowman:
  :return:
  """
  client = CatalogApiClient(Configuration(username=snowman.principal.client_id,
                                          password=snowman.credentials.client_secret.get_secret_value(),
                                          host=polaris_catalog_url))
  oauth_api = IcebergOAuth2API(client)
  token = oauth_api.get_token(scope='PRINCIPAL_ROLE:ALL', client_id=snowman.principal.client_id,
                              client_secret=snowman.credentials.client_secret.get_secret_value(),
                              grant_type='client_credentials',
                              _headers={'realm': 'POLARIS'})

  return IcebergCatalogAPI(CatalogApiClient(Configuration(access_token=token.access_token,
                                                          host=polaris_catalog_url)))


@pytest.fixture
def creator_catalog_client(polaris_catalog_url, creator):
  """
  Create an iceberg catalog client with TABLE_CREATE credentials
  :param polaris_catalog_url:
  :param creator:
  :return:
  """
  client = CatalogApiClient(Configuration(username=creator.principal.client_id,
                                          password=creator.credentials.client_secret.get_secret_value(),
                                          host=polaris_catalog_url))
  oauth_api = IcebergOAuth2API(client)
  token = oauth_api.get_token(scope='PRINCIPAL_ROLE:ALL', client_id=creator.principal.client_id,
                              client_secret=creator.credentials.client_secret.get_secret_value(),
                              grant_type='client_credentials',
                              _headers={'realm': 'POLARIS'})

  return IcebergCatalogAPI(CatalogApiClient(Configuration(access_token=token.access_token,
                                                          host=polaris_catalog_url)))


@pytest.fixture
def creator(polaris_url, polaris_catalog_url, root_client, snowflake_catalog):
  """
  create the creator principal with only TABLE_CREATE privileges
  :param root_client:
  :param snowflake_catalog:
  :return:
  """
  creator_name = "creator"
  principal_role = "creator_principal_role"
  catalog_role = "creator_catalog_role"
  try:
    creator = create_principal(polaris_url, polaris_catalog_url, root_client, creator_name)
    creator_principal_role = create_principal_role(root_client, principal_role)
    creator_catalog_role = create_catalog_role(root_client, snowflake_catalog, catalog_role)

    root_client.assign_catalog_role_to_principal_role(principal_role_name=creator_principal_role.name,
                                                      catalog_name=snowflake_catalog.name,
                                                      grant_catalog_role_request=GrantCatalogRoleRequest(
                                                        catalog_role=creator_catalog_role))
    root_client.add_grant_to_catalog_role(snowflake_catalog.name, creator_catalog_role.name,
                                          AddGrantRequest(grant=CatalogGrant(catalog_name=snowflake_catalog.name,
                                                                             type='catalog',
                                                                             privilege=CatalogPrivilege.TABLE_CREATE)))
    root_client.assign_principal_role(creator.principal.name,
                                      grant_principal_role_request=GrantPrincipalRoleRequest(
                                        principal_role=creator_principal_role))
    yield creator
  finally:
    root_client.delete_principal(creator_name)
    root_client.delete_principal_role(principal_role_name=principal_role)
    root_client.delete_catalog_role(catalog_role_name=catalog_role, catalog_name=snowflake_catalog.name)


@pytest.fixture
def reader_catalog_client(polaris_catalog_url, reader):
  """
  Create an iceberg catalog client with test_reader credentials
  :param polaris_catalog_url:
  :param reader:
  :return:
  """
  client = CatalogApiClient(Configuration(username=reader.principal.client_id,
                                          password=reader.credentials.client_secret.get_secret_value(),
                                          host=polaris_catalog_url))
  oauth_api = IcebergOAuth2API(client)
  token = oauth_api.get_token(scope='PRINCIPAL_ROLE:ALL', client_id=reader.principal.client_id,
                              client_secret=reader.credentials.client_secret.get_secret_value(),
                              grant_type='client_credentials',
                              _headers={'realm': 'POLARIS'})

  return IcebergCatalogAPI(CatalogApiClient(Configuration(access_token=token.access_token,
                                                          host=polaris_catalog_url)))


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_credentials(root_client, snowflake_catalog, polaris_catalog_url, snowman, reader):
  """
  Basic spark test - using snowman, create namespaces and a table. Insert into the table and read records back.

  Using the reader principal's credentials verify read access. Validate the reader cannot insert into the table.
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman:
  :param reader:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('SHOW NAMESPACES')
    spark.sql('USE db1.schema')
    spark.sql('CREATE TABLE iceberg_table (col1 int, col2 string)')
    spark.sql('SHOW TABLES')
    spark.sql("""INSERT INTO iceberg_table VALUES 
        (10, 'mystring'),
        (20, 'anotherstring'),
        (30, null)
        """)
    count = spark.sql("SELECT * FROM iceberg_table").count()
    assert count == 3

  # switch users to the reader. we can query, show namespaces, but we can't insert
  with IcebergSparkSession(credentials=f'{reader.principal.client_id}:{reader.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('SHOW NAMESPACES')
    spark.sql('USE db1.schema')
    count = spark.sql("SELECT * FROM iceberg_table").count()
    assert count == 3
    try:
      spark.sql("""INSERT INTO iceberg_table VALUES 
            (10, 'mystring'),
            (20, 'anotherstring'),
            (30, null)
            """)
      pytest.fail("Expected exception when trying to write without permission")
    except:
      print("Exception caught attempting to write without permission")

  # switch back to delete stuff
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('USE db1.schema')
    spark.sql('DROP TABLE iceberg_table PURGE')
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('DROP NAMESPACE db1.schema')
    spark.sql('DROP NAMESPACE db1')


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_cannot_create_table_outside_of_namespace_dir(root_client, snowflake_catalog, polaris_catalog_url,
                                                            snowman, reader):
  """
  Basic spark test - using snowman, create a namespace and try to create a table outside of the namespace. This should fail

  Using the reader principal's credentials verify read access. Validate the reader cannot insert into the table.
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman:
  :param reader:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    table_location = snowflake_catalog.properties.default_base_location + '/db1/outside_schema/table_outside_namespace'
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('SHOW NAMESPACES')
    spark.sql('USE db1.schema')
    try:
      spark.sql(f"CREATE TABLE iceberg_table_outside_namespace (col1 int, col2 string) LOCATION '{table_location}'")
      pytest.fail("Expected to fail when creating table outside of namespace directory")
    except Py4JJavaError as e:
      assert "is not in the list of allowed locations" in e.java_exception.getMessage()


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_creates_table_in_custom_namespace_dir(root_client, snowflake_catalog, polaris_catalog_url, snowman,
                                                     reader):
  """
  Basic spark test - using snowman, create a namespace and try to create a table outside of the namespace. This should fail

  Using the reader principal's credentials verify read access. Validate the reader cannot insert into the table.
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman:
  :param reader:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    namespace_location = snowflake_catalog.properties.default_base_location + '/db1/custom_location'
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql(f"CREATE NAMESPACE db1.schema LOCATION '{namespace_location}'")
    spark.sql('USE db1.schema')
    spark.sql("CREATE TABLE table_in_custom_namespace_location (col1 int, col2 string)")
    assert spark.sql("SELECT * FROM table_in_custom_namespace_location").count() == 0
    # check the metadata and assert the custom namespace location is used
    entries = spark.sql(
      "SELECT file FROM db1.schema.table_in_custom_namespace_location.metadata_log_entries").collect()
    assert namespace_location in entries[0][0]
    try:
        assert spark.sql("SELECT * FROM table_in_custom_namespace_location").count() == 0
        # check the metadata and assert the custom namespace location is used
        entries = spark.sql(
            "SELECT file FROM db1.schema.table_in_custom_namespace_location.metadata_log_entries").collect()
        assert namespace_location in entries[0][0]
    finally:
        spark.sql('DROP TABLE table_in_custom_namespace_location PURGE')


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_can_create_table_in_custom_allowed_dir(root_client, snowflake_catalog, polaris_catalog_url, snowman,
                                                      reader):
  """
  Basic spark test - using snowman, create a namespace and try to create a table outside of the namespace. This should fail

  Using the reader principal's credentials verify read access. Validate the reader cannot insert into the table.
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman:
  :param reader:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    table_location = snowflake_catalog.properties.default_base_location + '/db1/custom_schema_location/table_outside_namespace'
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql(
      f"CREATE NAMESPACE db1.schema LOCATION '{snowflake_catalog.properties.default_base_location}/db1/custom_schema_location'")
    spark.sql('SHOW NAMESPACES')
    spark.sql('USE db1.schema')
    # this is supported because it is inside of the custom namespace location
    spark.sql(f"CREATE TABLE iceberg_table_outside_namespace (col1 int, col2 string) LOCATION '{table_location}'")
    spark.sql("drop table iceberg_table_outside_namespace PURGE")


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_cannot_create_view_overlapping_table(root_client, snowflake_catalog, polaris_catalog_url, snowman,
                                                    reader):
  """
  Basic spark test - using snowman, create a namespace and try to create a table outside of the namespace. This should fail

  Using the reader principal's credentials verify read access. Validate the reader cannot insert into the table.
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman:
  :param reader:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    table_location = snowflake_catalog.properties.default_base_location + '/db1/schema/table_dir'
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql(f"CREATE NAMESPACE db1.schema LOCATION '{snowflake_catalog.properties.default_base_location}/db1/schema'")
    spark.sql('SHOW NAMESPACES')
    spark.sql('USE db1.schema')
    spark.sql(f"CREATE TABLE my_iceberg_table (col1 int, col2 string) LOCATION '{table_location}'")
    try:
      spark.sql(
        f"CREATE VIEW disallowed_view (int, string) TBLPROPERTIES ('location'= '{table_location}') AS SELECT * FROM my_iceberg_table")
      pytest.fail("Expected to fail when creating table outside of namespace directory")
    except Py4JJavaError as e:
      assert "conflicts with existing table or namespace at location" in e.java_exception.getMessage()
    spark.sql("drop table my_iceberg_table PURGE")


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_credentials_can_delete_after_purge(root_client, snowflake_catalog, polaris_catalog_url, snowman,
                                                  snowman_catalog_client, test_bucket, aws_bucket_base_location_prefix):
  """
  Using snowman, create namespaces and a table. Insert into the table in multiple operations and update existing records
  to generate multiple metadata.json files and manifests. Drop the table with purge=true. Poll S3 and validate all of
  the files are deleted.

  Using the reader principal's credentials verify read access. Validate the reader cannot insert into the table.
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman:
  :param reader:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    table_name = f'iceberg_test_table_{str(uuid.uuid4())[-10:]}'
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('SHOW NAMESPACES')
    spark.sql('USE db1.schema')
    spark.sql(f'CREATE TABLE {table_name} (col1 int, col2 string)')
    spark.sql('SHOW TABLES')

    # several inserts and an update, which should cause earlier files to show up as deleted in the later manifests
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (10, 'mystring'),
        (20, 'anotherstring'),
        (30, null)
        """)
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (40, 'mystring'),
        (50, 'anotherstring'),
        (60, null)
        """)
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (70, 'mystring'),
        (80, 'anotherstring'),
        (90, null)
        """)
    spark.sql(f"UPDATE {table_name} SET col2='changed string' WHERE col1 BETWEEN 20 AND 50")
    count = spark.sql(f"SELECT * FROM {table_name}").count()

    assert count == 9

    # fetch aws credentials to examine the metadata files
    response = snowman_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'), table_name,
                                                 "true")
    assert response.config is not None
    assert 's3.access-key-id' in response.config
    assert 's3.secret-access-key' in response.config
    assert 's3.session-token' in response.config

    s3 = boto3.client('s3',
                      aws_access_key_id=response.config['s3.access-key-id'],
                      aws_secret_access_key=response.config['s3.secret-access-key'],
                      aws_session_token=response.config['s3.session-token'])

    # Extract the table location from the metadata_location in the response
    # metadata_location format: s3://bucket/path/to/table/metadata/v1.metadata.json
    # We need to extract the base table path (everything before /metadata/)
    metadata_location = response.metadata_location
    assert metadata_location.startswith('s3://')
    # Remove s3:// prefix and bucket name to get the path
    path_without_scheme = metadata_location[5:]  # Remove 's3://'
    path_parts = path_without_scheme.split('/', 1)  # Split bucket and path
    bucket_from_metadata = path_parts[0]
    full_path = path_parts[1] if len(path_parts) > 1 else ''
    # Extract table base path (everything before /metadata/)
    table_base_path = full_path.rsplit('/metadata/', 1)[0] if '/metadata/' in full_path else ''

    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=f'{table_base_path}/data/')
    assert objects is not None
    assert 'Contents' in objects
    assert len(objects['Contents']) >= 4  # it varies - at least one file for each insert and one for the update
    print(f"Found {len(objects['Contents'])} data files in S3 before drop")

    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=f'{table_base_path}/metadata/')
    assert objects is not None
    assert 'Contents' in objects
    assert len(objects['Contents']) == 15  # 5 metadata.json files, 4 manifest lists, and 6 manifests
    print(f"Found {len(objects['Contents'])} metadata files in S3 before drop")

    # use the api client to ensure the purge flag is set to true
    snowman_catalog_client.drop_table(snowflake_catalog.name,
                                      codecs.decode("1F", "hex").decode("UTF-8").join(['db1', 'schema']), table_name,
                                      purge_requested=True)
    spark.sql('DROP NAMESPACE db1.schema')
    spark.sql('DROP NAMESPACE db1')
    print("Dropped table with purge - waiting for files to be deleted")
    attempts = 0

    # watch the data directory. metadata will be deleted first, so if data directory is clear, we can expect
    # metadata directory to be clear also
    while 'Contents' in objects and len(objects['Contents']) > 0 and attempts < 60:
      time.sleep(1)  # seconds, not milliseconds ;)
      objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                                Prefix=f'{table_base_path}/data/')
      attempts = attempts + 1

    if 'Contents' in objects and len(objects['Contents']) > 0:
      pytest.fail(f"Expected all data to be deleted, but found metadata files {objects['Contents']}")

    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=f'{table_base_path}/data/')
    if 'Contents' in objects and len(objects['Contents']) > 0:
      pytest.fail(f"Expected all data to be deleted, but found data files {objects['Contents']}")


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_credentials_can_write_with_random_prefix(root_client, snowflake_catalog, polaris_catalog_url, snowman,
                                                        snowman_catalog_client, test_bucket, aws_bucket_base_location_prefix):
  """
  Update the catalog configuration to support unstructured table locations. Using snowman, create namespaces and a
  table configured to use object-store layout in a folder under the catalog root, outside of the default table
  directory. Insert into the table in multiple operations and update existing records to generate multiple metadata.json
  files and manifests. Validate the data files are present under the expected subdirectory. Delete the files afterward.

  Using the reader principal's credentials verify read access. Validate the reader cannot insert into the table.
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman:
  :param reader:
  :return:
  """
  snowflake_catalog.properties.additional_properties['allow.unstructured.table.location'] = 'true'
  root_client.update_catalog(catalog_name=snowflake_catalog.name,
                             update_catalog_request=UpdateCatalogRequest(
                               properties=snowflake_catalog.properties.to_dict(),
                               current_entity_version=snowflake_catalog.entity_version))
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    table_name = f'iceberg_test_table_{str(uuid.uuid4())[-10:]}'
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('SHOW NAMESPACES')
    spark.sql('USE db1.schema')
    spark.sql(
      f"CREATE TABLE {table_name} (col1 int, col2 string) TBLPROPERTIES ('write.object-storage.enabled'='true','write.data.path'='s3://{test_bucket}/{aws_bucket_base_location_prefix}/snowflake_catalog/{table_name}data')")
    spark.sql('SHOW TABLES')

    # several inserts and an update, which should cause earlier files to show up as deleted in the later manifests
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (10, 'mystring'),
        (20, 'anotherstring'),
        (30, null)
        """)
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (40, 'mystring'),
        (50, 'anotherstring'),
        (60, null)
        """)
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (70, 'mystring'),
        (80, 'anotherstring'),
        (90, null)
        """)
    spark.sql(f"UPDATE {table_name} SET col2='changed string' WHERE col1 BETWEEN 20 AND 50")
    count = spark.sql(f"SELECT * FROM {table_name}").count()

    assert count == 9

    # fetch aws credentials to examine the metadata files
    response = snowman_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'), table_name,
                                                 "true")
    assert response.config is not None
    assert 's3.access-key-id' in response.config
    assert 's3.secret-access-key' in response.config
    assert 's3.session-token' in response.config

    s3 = boto3.client('s3',
                      aws_access_key_id=response.config['s3.access-key-id'],
                      aws_secret_access_key=response.config['s3.secret-access-key'],
                      aws_session_token=response.config['s3.session-token'])

    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/{table_name}data/')
    assert objects is not None
    assert len(objects['CommonPrefixes']) >= 3

    print(f"Found common prefixes in S3 {objects['CommonPrefixes']}")
    objs_to_delete = []

    for prefix in objects['CommonPrefixes']:
      # Don't utilize delimiter so all files (recursive) are returned
      data_objects = s3.list_objects(Bucket=test_bucket,
                            Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/{table_name}data/')
      assert data_objects is not None
      print(data_objects)
      assert 'Contents' in data_objects
      for obj in data_objects['Contents']:
          filePathMap = {'Key': obj['Key']}
          assert f'/schema/{table_name}/' in filePathMap['Key']
          objs_to_delete.append(filePathMap)

    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/{table_name}/metadata/')
    assert objects is not None
    assert 'Contents' in objects
    assert len(objects['Contents']) == 15  # 5 metadata.json files, 4 manifest lists, and 6 manifests
    print(f"Found {len(objects['Contents'])} metadata files in S3 before drop")

    # use the api client to ensure the purge flag is set to true
    snowman_catalog_client.drop_table(snowflake_catalog.name,
                                      codecs.decode("1F", "hex").decode("UTF-8").join(['db1', 'schema']), table_name,
                                      purge_requested=True)
    spark.sql('DROP NAMESPACE db1.schema')
    spark.sql('DROP NAMESPACE db1')
    s3.delete_objects(Bucket=test_bucket,
                      Delete={'Objects': objs_to_delete})


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_object_store_layout_under_table_dir(root_client, snowflake_catalog, polaris_catalog_url, snowman,
                                                   snowman_catalog_client, test_bucket, aws_bucket_base_location_prefix):
  """
  Using snowman, create namespaces and a table configured to use object-store layout, using a folder under the default
  table directory structure. Insert into the table in multiple operations and update existing records
  to generate multiple metadata.json files and manifests. Validate the data files are present under the expected
  subdirectory. Delete the files afterward.

  Using the reader principal's credentials verify read access. Validate the reader cannot insert into the table.
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman:
  :param reader:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    table_name = f'iceberg_test_table_{str(uuid.uuid4())[-10:]}'
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('SHOW NAMESPACES')
    spark.sql('USE db1.schema')
    table_base_dir = f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/{table_name}/obj_layout/'
    spark.sql(
      f"CREATE TABLE {table_name} (col1 int, col2 string) TBLPROPERTIES ('write.object-storage.enabled'='true','write.data.path'='s3://{test_bucket}/{table_base_dir}')")
    spark.sql('SHOW TABLES')

    # several inserts and an update, which should cause earlier files to show up as deleted in the later manifests
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (10, 'mystring'),
        (20, 'anotherstring'),
        (30, null)
        """)
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (40, 'mystring'),
        (50, 'anotherstring'),
        (60, null)
        """)
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (70, 'mystring'),
        (80, 'anotherstring'),
        (90, null)
        """)
    spark.sql(f"UPDATE {table_name} SET col2='changed string' WHERE col1 BETWEEN 20 AND 50")
    count = spark.sql(f"SELECT * FROM {table_name}").count()

    assert count == 9

    # fetch aws credentials to examine the metadata files
    response = snowman_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'), table_name,
                                                 "true")
    assert response.config is not None
    assert 's3.access-key-id' in response.config
    assert 's3.secret-access-key' in response.config
    assert 's3.session-token' in response.config

    s3 = boto3.client('s3',
                      aws_access_key_id=response.config['s3.access-key-id'],
                      aws_secret_access_key=response.config['s3.secret-access-key'],
                      aws_session_token=response.config['s3.session-token'])

    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=table_base_dir)
    assert objects is not None
    assert len(objects['CommonPrefixes']) >= 3

    print(f"Found common prefixes in S3 {objects['CommonPrefixes']}")
    objs_to_delete = []
    for prefix in objects['CommonPrefixes']:
      # Files may be under further subdirectories, so omit the
      # delimiter so that all files (recursive) are returned.
      data_objects = s3.list_objects(Bucket=test_bucket,
                                     Prefix=f'{prefix["Prefix"]}')
      assert data_objects is not None
      print(data_objects)
      assert 'Contents' in data_objects
      objs_to_delete.extend([{'Key': obj['Key']} for obj in data_objects['Contents']])

    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/{table_name}/metadata/')
    assert objects is not None
    assert 'Contents' in objects
    assert len(objects['Contents']) == 15  # 5 metadata.json files, 4 manifest lists, and 6 manifests
    print(f"Found {len(objects['Contents'])} metadata files in S3 before drop")

    # use the api client to ensure the purge flag is set to true
    snowman_catalog_client.drop_table(snowflake_catalog.name,
                                      codecs.decode("1F", "hex").decode("UTF-8").join(['db1', 'schema']), table_name,
                                      purge_requested=True)
    spark.sql('DROP NAMESPACE db1.schema')
    spark.sql('DROP NAMESPACE db1')
    s3.delete_objects(Bucket=test_bucket,
                      Delete={'Objects': objs_to_delete})


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
# @pytest.mark.skip(reason="This test is flaky")
def test_spark_credentials_can_create_views(snowflake_catalog, polaris_catalog_url, snowman):
  """
  Using snowman, create namespaces and a table. Insert into the table in multiple operations and update existing records
  to generate multiple metadata.json files and manifests. Create a view on the table. Verify the state of the view
  matches the state of the table.

  Using the reader principal's credentials verify read access. Validate the reader cannot insert into the table.
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    table_name = f'iceberg_test_table_{str(uuid.uuid4())[-10:]}'
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('SHOW NAMESPACES')
    spark.sql('USE db1.schema')
    spark.sql(f'CREATE TABLE {table_name} (col1 int, col2 string)')
    spark.sql('SHOW TABLES')

    # several inserts
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (10, 'mystring'),
        (20, 'anotherstring'),
        (30, null)
        """)
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (40, 'mystring'),
        (50, 'anotherstring'),
        (60, null)
        """)
    spark.sql(f"""INSERT INTO {table_name} VALUES 
        (70, 'mystring'),
        (80, 'anotherstring'),
        (90, null)
        """)
    # verify the view reflects the current state of the table
    spark.sql(f"CREATE VIEW {table_name}_view AS SELECT col2 FROM {table_name} where col1 > 30 ORDER BY col1 DESC")
    view_records = spark.sql(f"SELECT * FROM {table_name}_view").collect()
    assert len(view_records) == 6
    assert len(view_records[0]) == 1
    assert view_records[1][0] == 'anotherstring'
    assert view_records[5][0] == 'mystring'

    # Update some records. Assert the view reflects the new state
    spark.sql(f"UPDATE {table_name} SET col2='changed string' WHERE col1 BETWEEN 20 AND 50")
    view_records = spark.sql(f"SELECT * FROM {table_name}_view").collect()
    assert len(view_records) == 6
    assert view_records[5][0] == 'changed string'
    spark.sql(f"drop table {table_name} PURGE")


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_credentials_s3_direct_with_write(root_client, snowflake_catalog, polaris_catalog_url,
                                                snowman, snowman_catalog_client, test_bucket, aws_bucket_base_location_prefix):
  """
  Create two tables using Spark. Then call the loadTable api directly with snowman token to fetch the vended credentials
  for the first table.
  Verify that the credentials returned to snowman can read and write to the table's directory in S3, but don't allow
  reads or writes to the other table's directory
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman_catalog_client:
  :param reader_catalog_client:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('USE db1.schema')
    spark.sql('CREATE TABLE iceberg_table (col1 int, col2 string)')
    spark.sql('CREATE TABLE iceberg_table_2 (col1 int, col2 string)')

  table2_metadata = snowman_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'),
                                                      "iceberg_table_2",
                                                      "vended-credentials").metadata_location
  response = snowman_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'), "iceberg_table",
                                               "vended-credentials")
  assert response.config is not None
  assert 's3.access-key-id' in response.config
  assert 's3.secret-access-key' in response.config
  assert 's3.session-token' in response.config

  s3 = boto3.client('s3',
                    aws_access_key_id=response.config['s3.access-key-id'],
                    aws_secret_access_key=response.config['s3.secret-access-key'],
                    aws_session_token=response.config['s3.session-token'])

  objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                            Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/iceberg_table/')
  assert objects is not None
  assert 'CommonPrefixes' in objects
  assert len(objects['CommonPrefixes']) > 0

  objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                            Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/iceberg_table/metadata/')
  assert objects is not None
  assert 'Contents' in objects
  assert len(objects['Contents']) > 0

  metadata_file = next(f for f in objects['Contents'] if f['Key'].endswith('metadata.json'))
  assert metadata_file is not None

  metadata_contents = s3.get_object(Bucket=test_bucket, Key=metadata_file['Key'])
  assert metadata_contents is not None
  assert metadata_contents['ContentLength'] > 0

  put_object = s3.put_object(Bucket=test_bucket, Key=f"{metadata_file['Key']}.bak",
                             Body=metadata_contents['Body'].read())
  assert put_object is not None
  assert 'ETag' in put_object
  assert put_object['ETag'] is not None

  # list files in the other table's directory. The access policy should restrict this
  try:
    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/iceberg_table_2/metadata/')
    pytest.fail('Expected exception listing file outside of table directory')
  except botocore.exceptions.ClientError as error:
    print(error)

  try:
    metadata_contents = s3.get_object(Bucket=test_bucket, Key=table2_metadata)
    pytest.fail("Expected exception reading file outside of table directory")
  except botocore.exceptions.ClientError as error:
    print(error)

  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('USE db1.schema')
    spark.sql('DROP TABLE iceberg_table PURGE')
    spark.sql('DROP TABLE iceberg_table_2 PURGE')
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('DROP NAMESPACE db1.schema')
    spark.sql('DROP NAMESPACE db1')


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'false').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_credentials_s3_direct_without_write(root_client, snowflake_catalog, polaris_catalog_url,
                                                   snowman, reader_catalog_client, test_bucket, aws_bucket_base_location_prefix):
  """
  Create two tables using Spark. Then call the loadTable api directly with test_reader token to fetch the vended
  credentials for the first table.
  Verify that the credentials returned to test_reader allow reads, but don't allow writes to the table's directory
  and don't allow reads or writes anywhere else on S3. This verifies that Polaris's authz model does not only prevent
  users from updating metadata to enforce read-only access, but uses credential scoping to enforce restrictions at
  the storage layer.
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param reader_catalog_client:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('USE db1.schema')
    spark.sql('CREATE TABLE iceberg_table (col1 int, col2 string)')
    spark.sql('CREATE TABLE iceberg_table_2 (col1 int, col2 string)')

  table2_metadata = reader_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'),
                                                     "iceberg_table_2",
                                                     "vended-credentials").metadata_location

  response = reader_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'), "iceberg_table",
                                              "vended-credentials")
  assert response.config is not None
  assert 's3.access-key-id' in response.config
  assert 's3.secret-access-key' in response.config
  assert 's3.session-token' in response.config

  s3 = boto3.client('s3',
                    aws_access_key_id=response.config['s3.access-key-id'],
                    aws_secret_access_key=response.config['s3.secret-access-key'],
                    aws_session_token=response.config['s3.session-token'])

  objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                            Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/iceberg_table/metadata/')
  assert objects is not None
  assert 'Contents' in objects
  assert len(objects['Contents']) > 0

  metadata_file = next(f for f in objects['Contents'] if f['Key'].endswith('metadata.json'))
  assert metadata_file is not None

  metadata_contents = s3.get_object(Bucket=test_bucket, Key=metadata_file['Key'])
  assert metadata_contents is not None
  assert metadata_contents['ContentLength'] > 0

  # try to write. Expect it to fail
  try:
    put_object = s3.put_object(Bucket=test_bucket, Key=f"{metadata_file['Key']}.bak",
                               Body=metadata_contents['Body'].read())
    pytest.fail("Expect exception trying to write to table directory")
  except botocore.exceptions.ClientError as error:
    print(error)

  # list files in the other table's directory. The access policy should restrict this
  try:
    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/iceberg_table_2/metadata/')
    pytest.fail('Expected exception listing file outside of table directory')
  except botocore.exceptions.ClientError as error:
    print(error)

  try:
    metadata_contents = s3.get_object(Bucket=test_bucket, Key=table2_metadata)
    pytest.fail("Expected exception reading file outside of table directory")
  except botocore.exceptions.ClientError as error:
    print(error)

  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('USE db1.schema')
    spark.sql('DROP TABLE iceberg_table')
    spark.sql('DROP TABLE iceberg_table_2')
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('DROP NAMESPACE db1.schema')
    spark.sql('DROP NAMESPACE db1')


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'false').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_credentials_s3_direct_without_read(
    snowflake_catalog, snowman_catalog_client, creator_catalog_client, test_bucket):
  """
  Create a table using `creator`, which does not have TABLE_READ_DATA and expect a `ForbiddenException`
  """
  snowman_catalog_client.create_namespace(
    prefix=snowflake_catalog.name,
    create_namespace_request=CreateNamespaceRequest(
      namespace=["some_schema"]
    )
  )

  try:
    creator_catalog_client.create_table(
      prefix=snowflake_catalog.name,
      namespace="some_schema",
      x_iceberg_access_delegation="true",
      create_table_request=CreateTableRequest(
        name="some_table",
        var_schema=ModelSchema(
          type='struct',
          fields=[],
        )
      )
    )
    pytest.fail("Expected exception when creating a table without TABLE_WRITE")
  except Exception as e:
    assert 'CREATE_TABLE_DIRECT_WITH_WRITE_DELEGATION' in str(e)

  snowman_catalog_client.drop_namespace(
    prefix=snowflake_catalog.name,
    namespace="some_schema"
  )


def create_principal(polaris_url, polaris_catalog_url, api, principal_name):
  principal = Principal(name=principal_name, type="SERVICE")
  principal_result = api.create_principal(CreatePrincipalRequest(principal=principal))

  token_client = CatalogApiClient(Configuration(username=principal_result.principal.client_id,
                                                  password=principal_result.credentials.client_secret.get_secret_value(),
                                                  host=polaris_catalog_url))
  oauth_api = IcebergOAuth2API(token_client)
  token = oauth_api.get_token(scope='PRINCIPAL_ROLE:ALL', client_id=principal_result.principal.client_id,
                                client_secret=principal_result.credentials.client_secret.get_secret_value(),
                                grant_type='client_credentials',
                                _headers={'realm': 'POLARIS'})
  rotate_client = ManagementApiClient(Configuration(access_token=token.access_token,
                                                      host=polaris_url))
  rotate_api = PolarisDefaultApi(rotate_client)

  rotate_credentials = rotate_api.rotate_credentials(principal_name=principal_name)
  return rotate_credentials


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_credentials_s3_scoped_to_metadata_data_locations(root_client, snowflake_catalog, polaris_catalog_url,
                                                                snowman, snowman_catalog_client, test_bucket, aws_bucket_base_location_prefix):
  """
  Create a table using Spark. Then call the loadTable api directly with snowman token to fetch the vended credentials
  for the table.
  Verify that the credentials returned to snowman can only work for the location that ending with metadata or data directory
  :param root_client:
  :param snowflake_catalog:
  :param polaris_catalog_url:
  :param snowman_catalog_client:
  :param reader_catalog_client:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('USE db1.schema')
    spark.sql('CREATE TABLE iceberg_table_scope_loc(col1 int, col2 string)')
    spark.sql(
      f'''CREATE TABLE iceberg_table_scope_loc_slashes (col1 int, col2 string) LOCATION \'s3://{test_bucket}/{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/iceberg_table_scope_loc_slashes/path_with_slashes///////\'''')

  prefix1 = f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/iceberg_table_scope_loc'
  prefix2 = f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/iceberg_table_scope_loc_slashes/path_with_slashes'
  response1 = snowman_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'),
                                                "iceberg_table_scope_loc",
                                                "vended-credentials")
  response2 = snowman_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'),
                                                "iceberg_table_scope_loc_slashes",
                                                "vended-credentials")
  assert response1 is not None
  assert response2 is not None
  assert response1.metadata_location.startswith(f"s3://{test_bucket}/{prefix1}/metadata/")
  # ensure that the slashes are removed before "/metadata/"
  assert response2.metadata_location.startswith(f"s3://{test_bucket}/{prefix2}/metadata/")

  s3_1 = boto3.client('s3',
                      aws_access_key_id=response1.config['s3.access-key-id'],
                      aws_secret_access_key=response1.config['s3.secret-access-key'],
                      aws_session_token=response1.config['s3.session-token'])

  s3_2 = boto3.client('s3',
                      aws_access_key_id=response2.config['s3.access-key-id'],
                      aws_secret_access_key=response2.config['s3.secret-access-key'],
                      aws_session_token=response2.config['s3.session-token'])
  for client, prefix in [(s3_1, prefix1), (s3_2, prefix2)]:
    objects = client.list_objects(Bucket=test_bucket, Delimiter='/',
                                  Prefix=f'{prefix}/metadata/')
    assert objects is not None
    assert 'Contents' in objects, f'list metadata files failed in prefix: {prefix}/metadata/'

    objects = client.list_objects(Bucket=test_bucket, Delimiter='/',
                                  Prefix=f'{prefix}/data/')
    assert objects is not None
    # no insert executed, so should not have any data files
    assert 'Contents' not in objects, f'No contents should be in prefix: {prefix}/data/'

    objects = client.list_objects(Bucket=test_bucket, Delimiter='/',
                                  Prefix=f'{prefix}/')
    assert objects is not None
    assert 'CommonPrefixes' in objects, f'list prefixes failed in prefix: {prefix}/'
    assert len(objects['CommonPrefixes']) > 0

  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('USE db1.schema')
    spark.sql('DROP TABLE iceberg_table_scope_loc PURGE')
    spark.sql('DROP TABLE iceberg_table_scope_loc_slashes PURGE')
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('DROP NAMESPACE db1.schema')
    spark.sql('DROP NAMESPACE db1')


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_ctas(snowflake_catalog, polaris_catalog_url, snowman):
  """
  Create a table using CTAS and ensure that credentials are vended
  :param root_client:
  :param snowflake_catalog:
  :return:
  """
  with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                           catalog_name=snowflake_catalog.name,
                           polaris_url=polaris_catalog_url) as spark:
    table_name = f'iceberg_test_table_{str(uuid.uuid4())[-10:]}'
    spark.sql(f'USE {snowflake_catalog.name}')
    spark.sql('CREATE NAMESPACE db1')
    spark.sql('CREATE NAMESPACE db1.schema')
    spark.sql('USE db1.schema')
    spark.sql(f'CREATE TABLE {table_name}_t1 (col1 int)')
    spark.sql('SHOW TABLES')

    # Insert some data
    spark.sql(f"INSERT INTO {table_name}_t1 VALUES (10)")

    # Run CTAS
    spark.sql(f"CREATE TABLE {table_name}_t2 AS SELECT * FROM {table_name}_t1")
    spark.sql(f"drop table {table_name}_t1 PURGE")
    spark.sql(f"drop table {table_name}_t2 PURGE")


@pytest.mark.skipif(os.environ.get('AWS_TEST_ENABLED', 'False').lower() != 'true',
                    reason='AWS_TEST_ENABLED is not set or is false')
def test_spark_credentials_s3_exception_on_metadata_file_deletion(root_client, snowflake_catalog, polaris_catalog_url,
                                                snowman, snowman_catalog_client, test_bucket, aws_bucket_base_location_prefix):
    """
    Create a using Spark. Then call the loadTable api directly with snowman token to fetch the vended credentials
    for the first table.
    Delete the metadata directory and try to access the table using the vended credentials.
    It should throw 404 exception
    :param root_client:
    :param snowflake_catalog:
    :param polaris_catalog_url:
    :param snowman_catalog_client:
    :param reader_catalog_client:
    :return:
    """
    with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                             catalog_name=snowflake_catalog.name,
                             polaris_url=polaris_catalog_url) as spark:
        spark.sql(f'USE {snowflake_catalog.name}')
        spark.sql('CREATE NAMESPACE db1')
        spark.sql('CREATE NAMESPACE db1.schema')
        spark.sql('USE db1.schema')
        spark.sql('CREATE TABLE iceberg_table (col1 int, col2 string)')

    response = snowman_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'),
                                                        "iceberg_table",
                                                        "vended-credentials")
    assert response.config is not None
    assert 's3.access-key-id' in response.config
    assert 's3.secret-access-key' in response.config
    assert 's3.session-token' in response.config

    s3 = boto3.client('s3',
                      aws_access_key_id=response.config['s3.access-key-id'],
                      aws_secret_access_key=response.config['s3.secret-access-key'],
                      aws_session_token=response.config['s3.session-token'])

    # Get metadata files
    objects = s3.list_objects(Bucket=test_bucket, Delimiter='/',
                              Prefix=f'{aws_bucket_base_location_prefix}/snowflake_catalog/db1/schema/iceberg_table/metadata/')
    assert objects is not None
    assert 'Contents' in objects
    assert len(objects['Contents']) > 0

    # Verify metadata content
    metadata_file = next(f for f in objects['Contents'] if f['Key'].endswith('metadata.json'))
    assert metadata_file is not None

    metadata_contents = s3.get_object(Bucket=test_bucket, Key=metadata_file['Key'])
    assert metadata_contents is not None
    assert metadata_contents['ContentLength'] > 0

    # Delete metadata files
    objects_to_delete = [{'Key': obj['Key']} for obj in objects['Contents']]
    s3.delete_objects(Bucket=test_bucket,
                      Delete={'Objects': objects_to_delete})

    try:
        response = snowman_catalog_client.load_table(snowflake_catalog.name, unquote('db1%1Fschema'),
                                                     "iceberg_table",
                                                     "vended-credentials")
    except Exception as e:
        # 400 error(BadRequest) is thrown when metadata file is missing
        assert '400' in str(e)


    with IcebergSparkSession(credentials=f'{snowman.principal.client_id}:{snowman.credentials.client_secret.get_secret_value()}',
                             catalog_name=snowflake_catalog.name,
                             polaris_url=polaris_catalog_url) as spark:
        # Spark drop table triggers load table underneath, which fails due to missing metadata file.
        # Directly call drop_table api to drop the table entity.
        snowman_catalog_client.drop_table(snowflake_catalog.name, unquote('db1%1Fschema'),
                            "iceberg_table")
        spark.sql(f'USE {snowflake_catalog.name}')
        spark.sql('DROP NAMESPACE db1.schema')
        spark.sql('DROP NAMESPACE db1')

def create_catalog_role(api, catalog, role_name):
  catalog_role = CatalogRole(name=role_name)
  try:
    api.create_catalog_role(catalog_name=catalog.name,
                            create_catalog_role_request=CreateCatalogRoleRequest(catalog_role=catalog_role))
    return api.get_catalog_role(catalog_name=catalog.name, catalog_role_name=role_name)
  except ApiException:
    return api.get_catalog_role(catalog_name=catalog.name, catalog_role_name=role_name)
  else:
    raise e


def create_principal_role(api, role_name):
  principal_role = PrincipalRole(name=role_name)
  try:
    api.create_principal_role(CreatePrincipalRoleRequest(principal_role=principal_role))
    return api.get_principal_role(principal_role_name=role_name)
  except ApiException:
    return api.get_principal_role(principal_role_name=role_name)
