Skip to content

Commit

Permalink
fix: revise parameter added for unity read/write mode
Browse files Browse the repository at this point in the history
  • Loading branch information
pmogren committed Feb 14, 2025
1 parent 30125d0 commit 2c4495e
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 5 deletions.
2 changes: 1 addition & 1 deletion CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ When proposing features, please include:

To set up your development environment:

1. Ensure that your system has a suitable Python version installed (>=3.7, <=3.11)
1. Ensure that your system has a suitable Python version installed (>=3.9, <=3.11)
2. [Install the Rust compilation toolchain](https://www.rust-lang.org/tools/install)
3. Clone the Daft repo: `git clone git@github.com:Eventual-Inc/Daft.git`
4. Run `make .venv` from your new cloned Daft repository to create a new virtual environment with all of Daft's development dependencies installed
Expand Down
9 changes: 5 additions & 4 deletions daft/unity_catalog/unity_catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import dataclasses
import warnings
from typing import Callable
from typing import Callable, Literal
from urllib.parse import urlparse

import unitycatalog
Expand Down Expand Up @@ -91,14 +91,14 @@ def _paginated_list_tables(client: unitycatalog.Unitycatalog, page_token: str |
return self._paginate_to_completion(_paginated_list_tables)

def load_table(
self, table_name: str, new_table_storage_path: str | None = None, intent: str = "READ_WRITE"
self, table_name: str, new_table_storage_path: str | None = None, mode: Literal["r"|"rw"] = "rw"
) -> UnityCatalogTable:
"""Loads an existing Unity Catalog table. If the table is not found, and information is provided in the method to create a new table, a new table will be attempted to be registered.
Args:
table_name (str): Name of the table in Unity Catalog in the form of dot-separated, 3-level namespace
new_table_storage_path (str, optional): Cloud storage path URI to register a new external table using this path. Unity Catalog will validate if the path is valid and authorized for the principal, else will raise an exception.
intent (str, optional): The intended use of the table, which impacts authorization. Defaults to "READ_WRITE", else can be "READ".
mode (Literal["r", "rw"], optional): The intended use of the table, which impacts authorization. Defaults to "rw".
Returns:
UnityCatalogTable
Expand Down Expand Up @@ -140,7 +140,8 @@ def load_table(
table_id = table_info.table_id
storage_location = table_info.storage_location
# Grab credentials from Unity catalog and place it into the Table
temp_table_credentials = self._client.temporary_table_credentials.create(operation=intent, table_id=table_id)
operation = "READ" if mode == "r" else "READ_WRITE"
temp_table_credentials = self._client.temporary_table_credentials.create(operation=operation, table_id=table_id)

scheme = urlparse(storage_location).scheme
if scheme == "s3" or scheme == "s3a":
Expand Down

0 comments on commit 2c4495e

Please sign in to comment.