Kevinjqliu/adls token

[kevinjqliu]

Rationale for this change

Are these changes tested?

Are there any user-facing changes?

[kevinjqliu]

/gemini review

[gemini-code-assist]

Summary of Changes

Hello @kevinjqliu, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a suite of enhancements to PyIceberg, primarily focusing on expanding authentication capabilities for REST catalogs and cloud storage, refining table maintenance operations, and bolstering schema evolution and column projection logic. Key updates include the addition of Google authentication for REST catalogs, improved ADLS token and credential handling, and a new, more intuitive API for managing table snapshots. Furthermore, the changes significantly improve how PyIceberg processes schemas and data, especially concerning missing columns and nanosecond timestamp precision, ensuring greater compatibility and data integrity across various Iceberg table versions and data sources.

Highlights

• Google Authentication for REST Catalog: The REST catalog now supports Google authentication, allowing users to configure their catalog to authenticate using Google Application Default Credentials or a specified service account key file. This significantly enhances security and integration with Google Cloud services.
• Enhanced ADLS Authentication: Expanded Azure Data Lake Storage (ADLS) authentication options now include direct token and credential support, as well as authentication via client ID, client secret, and tenant ID. This provides greater flexibility and security for connecting to ADLS Gen2.
• Refactored Table Maintenance API: The table maintenance operations, particularly snapshot expiration, have been refactored into a new table.maintenance API. This provides a more intuitive and organized interface for managing table lifecycle, allowing for operations like expiring snapshots by ID or by age.
• Improved Schema Evolution and Column Projection: Significant improvements have been made to how PyIceberg handles schema evolution and column projection. The system can now correctly evaluate expressions against columns that might be missing in the physical data file by leveraging projected values or default field values, ensuring data consistency across schema changes. Additionally, PyArrow schema conversion now properly accounts for Iceberg table format versions, including support for nanosecond timestamp precision in format version 3.
• Robust Avro Field Name Sanitization: To ensure broader compatibility and correct data representation, Avro field names are now comprehensively sanitized. This includes handling special characters and emojis, with the original field name preserved using a new iceberg-field-name property in the Avro schema.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces significant enhancements and refactoring across the PyIceberg library. Key additions include authentication support for ADLS using tokens and for the REST catalog using Google credentials. The table maintenance API has been improved by moving snapshot expiration into a dedicated maintenance submodule, which clarifies the API structure. The implementation of column projection for filters is a major feature that will improve query performance. There are also important updates for nanosecond timestamp precision and Avro schema sanitization to improve cross-platform compatibility. The codebase has been improved with various refactorings and bug fixes, and the test suite has been substantially expanded to cover the new functionality.

[gemini-code-assist]

To improve readability and maintainability, it's better to avoid magic numbers. Please consider replacing 3600 with a named constant like ONE_HOUR_IN_SECONDS.

Suggested change

	self._expires_on = int(time.time()) + 3600
	self._expires_on = int(time.time()) + (60 * 60)

[gemini-code-assist]

This block of code for determining the credential can be simplified for better readability and conciseness. Using an assignment expression (walrus operator) can make the code more direct.

    credential = None
    if token := properties.get(ADLS_TOKEN):
        credential = StaticTokenCredential(token)
    elif ADLS_CREDENTIAL in properties:
        credential = properties.get(ADLS_CREDENTIAL)

[gemini-code-assist]

The logic for validating that all required credential keys are provided can be expressed more concisely and efficiently using set operations. This would also make the intent of the code clearer.

Suggested change

	# Validate that all three are provided together for ClientSecretCredential
	credential_keys = ["client_id", "client_secret", "tenant_id"]
	provided_keys = [key for key in credential_keys if key in client_kwargs]
	if provided_keys and len(provided_keys) != len(credential_keys):
	missing_keys = [key for key in credential_keys if key not in client_kwargs]
	raise ValueError(
	f"client_id, client_secret, and tenant_id must all be provided together "
	f"to use ClientSecretCredential for Azure authentication. "
	f"Provided: {provided_keys}, Missing: {missing_keys}"
	)
	# Validate that all three are provided together for ClientSecretCredential
	credential_keys = {"client_id", "client_secret", "tenant_id"}
	provided_keys = credential_keys.intersection(client_kwargs)
	if provided_keys and provided_keys != credential_keys:
	missing_keys = credential_keys - provided_keys
	raise ValueError(
	f"client_id, client_secret, and tenant_id must all be provided together "
	f"to use ClientSecretCredential for Azure authentication. "
	f"Provided: {sorted(list(provided_keys))}, Missing: {sorted(list(missing_keys))}"
	)

[gemini-code-assist]

Code Review

This pull request introduces a wide range of improvements and new features across the PyIceberg library. Key additions include support for Google Cloud and enhanced ADLS authentication, providing more flexibility for users on different cloud platforms. The table maintenance API has been refactored, with operations like snapshot expiration now accessible through a cleaner table.maintenance interface. There are also significant enhancements to column projection and initial support for nanosecond precision timestamps, aligning with the latest Iceberg specifications. The transaction management has been made more robust. Overall, these are excellent changes that improve the library's capabilities, API design, and correctness. I have a couple of suggestions for improvement.

[gemini-code-assist]

This logic for column projection has a couple of potential issues. First, it's missing IndexError handling for schemas that have evolved, which could cause crashes when processing older data files that don't have a new partition field. Second, the if partition_value := ... check will incorrectly skip falsy but valid partition values (like 0 or False), leading to incorrect data projection.

I suggest restoring the try...except block and ensuring all non-missing values are projected, regardless of their truthiness.

Suggested change

	if partition_value := accessors[partition_field.field_id].get(file.partition):
	projected_missing_fields[field_id] = partition_value
	try:
	projected_missing_fields[field_id] = accessors[partition_field.field_id].get(file.partition)
	except IndexError:
	# This can happen when a new partition field was added after the file was written.
	# In this case, we can't project the value.
	pass

[gemini-code-assist]

The StaticTokenCredential hardcodes the token expiry to one hour. This could be problematic if the provided token has a shorter lifetime. It would be more robust to allow configuring the expiry time, or at least clearly document this fixed expiry behavior in the class's docstring.

    class StaticTokenCredential(AsyncTokenCredential):
        """A credential that uses a static token.

        The token expiry is set to 1 hour from initialization.
        """
        def __init__(self, token_string: str) -> None:
            self._token = token_string
            self._expires_on = int(time.time()) + 3600

        async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
            return AccessToken(self._token, self._expires_on)