deltacat 1.1.38__py3-none-any.whl → 2.0.0__py3-none-any.whl

deltacat/compute/converter/utils/s3u.py

@@ -0,0 +1,148 @@
+from tenacity import (
+    Retrying,
+    retry_if_exception_type,
+    stop_after_delay,
+    wait_random_exponential,
+)
+from typing import Union, Optional, Dict, Any, List, Callable
+from deltacat.types.tables import (
+    CapturedBlockWritePaths,
+    UuidBlockWritePathProvider,
+)
+from deltacat.types.tables import (
+    get_table_writer,
+    get_table_length,
+    TABLE_CLASS_TO_SLICER_FUNC,
+)
+from deltacat.exceptions import RetryableError
+from deltacat.storage import (
+    DistributedDataset,
+    LocalTable,
+)
+from deltacat.types.media import (
+    ContentEncoding,
+    ContentType,
+)
+from deltacat.constants import UPLOAD_SLICED_TABLE_RETRY_STOP_AFTER_DELAY
+import s3fs
+import boto3
+from boto3.session import Session
+from botocore.credentials import Credentials
+
+
+def get_credential() -> Credentials:
+    boto3_session: Session = boto3.Session()
+    credentials: Credentials = boto3_session.get_credentials()
+    return credentials
+
+
+def get_s3_file_system(content_type):
+    import s3fs  # noqa: F401
+
+    token_holder = get_credential()
+    content_encoding = ContentEncoding.IDENTITY
+
+    s3_file_system = s3fs.S3FileSystem(
+        key=token_holder.access_key,
+        secret=token_holder.secret_key,
+        token=token_holder.token,
+        s3_additional_kwargs={
+            "ServerSideEncryption": "aws:kms",
+            # TODO: Get tagging from table properties
+            "ContentType": content_type.value,
+            "ContentEncoding": content_encoding.value,
+        },
+    )
+    return s3_file_system
+
+
+def upload_table_with_retry(
+    table: Union[LocalTable, DistributedDataset],
+    s3_url_prefix: str,
+    s3_table_writer_kwargs: Optional[Dict[str, Any]],
+    content_type: ContentType = ContentType.PARQUET,
+    max_records_per_file: Optional[int] = 4000000,
+    filesystem: Optional[s3fs.S3FileSystem] = None,
+    **s3_client_kwargs: Any,
+) -> List[str]:
+    """
+    Writes the given table to 1 or more S3 files and return the paths
+    of the S3 files written.
+    """
+    retrying = Retrying(
+        wait=wait_random_exponential(multiplier=1, max=60),
+        stop=stop_after_delay(UPLOAD_SLICED_TABLE_RETRY_STOP_AFTER_DELAY),
+        retry=retry_if_exception_type(RetryableError),
+    )
+
+    if s3_table_writer_kwargs is None:
+        s3_table_writer_kwargs = {}
+
+    if not filesystem:
+        filesystem = get_s3_file_system(content_type=content_type)
+    capture_object = CapturedBlockWritePaths()
+    block_write_path_provider = UuidBlockWritePathProvider(
+        capture_object=capture_object, base_path=s3_url_prefix
+    )
+    s3_table_writer_func = get_table_writer(table)
+    table_record_count = get_table_length(table)
+    if max_records_per_file is None or not table_record_count:
+        retrying(
+            fn=upload_table,
+            table_slices=table,
+            s3_base_url=f"{s3_url_prefix}",
+            s3_file_system=filesystem,
+            s3_table_writer_func=s3_table_writer_func,
+            s3_table_writer_kwargs=s3_table_writer_kwargs,
+            block_write_path_provider=block_write_path_provider,
+            content_type=content_type,
+            **s3_client_kwargs,
+        )
+    else:
+        table_slicer_func = TABLE_CLASS_TO_SLICER_FUNC.get(type(table))
+        table_slices = table_slicer_func(table, max_records_per_file)
+        for table_slice in table_slices:
+            retrying(
+                fn=upload_table,
+                table_slices=table_slice,
+                s3_base_url=f"{s3_url_prefix}",
+                s3_file_system=filesystem,
+                s3_table_writer_func=s3_table_writer_func,
+                s3_table_writer_kwargs=s3_table_writer_kwargs,
+                block_write_path_provider=block_write_path_provider,
+                content_type=content_type,
+                **s3_client_kwargs,
+            )
+    del block_write_path_provider
+    write_paths = capture_object.write_paths()
+    s3_write_paths = []
+    for path in write_paths:
+        s3_write_path = construct_s3_url(path)
+        s3_write_paths.append(s3_write_path)
+    return s3_write_paths
+
+
+def construct_s3_url(path: Optional[str]) -> Optional[str]:
+    if path:
+        return f"s3://{path}"
+    return None
+
+
+def upload_table(
+    table_slices: Union[LocalTable, DistributedDataset],
+    s3_base_url: str,
+    s3_file_system: s3fs.S3FileSystem,
+    s3_table_writer_func: Callable,
+    block_write_path_provider: UuidBlockWritePathProvider,
+    content_type: ContentType,
+    s3_table_writer_kwargs: Dict[str, Any],
+) -> None:
+    s3_table_writer_func(
+        table_slices,
+        s3_base_url,
+        s3_file_system,
+        block_write_path_provider,
+        content_type.value,
+        **s3_table_writer_kwargs,
+    )
+    # TODO: Add a proper fix for block_refs and write_paths not persisting in Ray actors

deltacat/compute/janitor.py

@@ -0,0 +1,205 @@
+import time
+import os
+import posixpath
+import pyarrow.fs
+from pyarrow.fs import FileSelector, FileType
+from itertools import chain
+from deltacat.storage.model.transaction import Transaction
+from deltacat.utils.filesystem import resolve_path_and_filesystem
+from deltacat.constants import (
+    TXN_DIR_NAME,
+    RUNNING_TXN_DIR_NAME,
+    FAILED_TXN_DIR_NAME,
+    TXN_PART_SEPARATOR,
+)
+from deltacat.storage.model.types import TransactionState
+import logging
+from deltacat import logs
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def brute_force_search_matching_metafiles(
+    dirty_files_names, filesystem: pyarrow.fs.FileSystem, catalog_root
+):
+    txn_dir_name = TXN_DIR_NAME
+    # collect transaction ids of the files
+    transaction_ids = []
+    for dirty_file in dirty_files_names:
+        parts = dirty_file.split(TXN_PART_SEPARATOR)
+        if len(parts) < 2:
+            continue
+        transaction_ids.append(parts[1])
+
+    def recursive_search(path):
+        try:
+            selector = FileSelector(path, recursive=False)
+            entries = filesystem.get_file_info(selector)
+        except Exception as e:
+            logger.error(f"Error listing directory '{path}': {e}")
+            return
+
+        for entry in entries:
+            base_name = posixpath.basename(entry.path)
+            if entry.type == FileType.File:
+                for transaction_id in transaction_ids:
+                    # Look for transaction_id in the filename
+                    if transaction_id in base_name:
+                        try:
+                            filesystem.delete_file(entry.path)
+                            logger.debug(f"Deleted file: {entry.path}")
+                        except Exception as e:
+                            logger.error(f"Error deleting file '{entry.path}': {e}")
+
+            elif entry.type == FileType.Directory:
+                # Skip directories that match txn_dir_name
+                if posixpath.basename(entry.path) == txn_dir_name:
+                    logger.debug(f"Skipping directory: {entry.path}")
+                    continue
+                recursive_search(entry.path)
+
+    # Start recursive search from the catalog root
+    recursive_search(catalog_root)
+
+    # renaming to successful completion
+    for dirty_file in dirty_files_names:
+        failed_txn_log_dir = posixpath.join(
+            catalog_root, TXN_DIR_NAME, FAILED_TXN_DIR_NAME
+        )
+        old_log_path = posixpath.join(failed_txn_log_dir, dirty_file)
+
+        # new_filename = dirty_file.replace(TIMEOUT_TXN, SUCCESSFULLY_CLEANED)
+        new_log_path = posixpath.join(failed_txn_log_dir, dirty_file)
+        try:
+            filesystem.move(old_log_path, new_log_path)
+            logger.debug(f"Renamed file from {old_log_path} to {new_log_path}")
+        except Exception as e:
+            logger.error(f"Error renaming file '{old_log_path}': {e}")
+
+
+def janitor_delete_timed_out_transaction(catalog_root: str) -> None:
+    """
+    Traverse the running transactions directory and move transactions that have been
+    running longer than the threshold into the failed transactions directory.
+    """
+    catalog_root_normalized, filesystem = resolve_path_and_filesystem(catalog_root)
+
+    txn_log_dir = posixpath.join(catalog_root_normalized, TXN_DIR_NAME)
+    running_txn_log_dir = posixpath.join(txn_log_dir, RUNNING_TXN_DIR_NAME)
+    failed_txn_log_dir = posixpath.join(txn_log_dir, FAILED_TXN_DIR_NAME)
+
+    dirty_files = []
+
+    running_txn_file_selector = FileSelector(running_txn_log_dir, recursive=False)
+    running_txn_info_list = filesystem.get_file_info(running_txn_file_selector)
+
+    for running_txn_info in running_txn_info_list:
+        try:
+            filename = posixpath.basename(running_txn_info.path)
+            parts = filename.split(TXN_PART_SEPARATOR)
+            end_time_str = parts[-1]
+            end_time = float(end_time_str)
+            current_time = time.time_ns()
+            if end_time <= current_time:
+                src_path = running_txn_info.path
+                new_filename = f"{filename}"
+                dest_path = posixpath.join(failed_txn_log_dir, new_filename)
+
+                # Move the file using copy and delete
+                with filesystem.open_input_file(src_path) as src_file:
+                    contents = src_file.read()
+
+                with filesystem.open_output_stream(dest_path) as dest_file:
+                    dest_file.write(contents)
+                filesystem.delete_file(src_path)
+
+                dirty_files.append(new_filename)
+
+        except Exception as e:
+            logger.error(
+                f"Error cleaning failed transaction '{running_txn_info.path}': {e}"
+            )
+
+    # Pass catalog_root to the brute force search so it searches from the right place
+    brute_force_search_matching_metafiles(
+        dirty_files, filesystem, catalog_root_normalized
+    )
+
+
+def janitor_remove_files_in_failed(
+    catalog_root: str, filesystem: pyarrow.fs.FileSystem = None
+) -> None:
+    """
+    Cleans up metafiles and locator files associated with failed transactions.
+    """
+    if filesystem is None:
+        catalog_root_normalized, filesystem = resolve_path_and_filesystem(catalog_root)
+    else:
+        catalog_root_normalized, filesystem = resolve_path_and_filesystem(
+            catalog_root, filesystem
+        )
+
+    txn_log_dir = posixpath.join(catalog_root_normalized, TXN_DIR_NAME)
+    failed_txn_log_dir = posixpath.join(txn_log_dir, FAILED_TXN_DIR_NAME)
+    running_txn_log_dir = posixpath.join(txn_log_dir, RUNNING_TXN_DIR_NAME)
+    filesystem.create_dir(failed_txn_log_dir, recursive=True)
+
+    failed_txn_file_selector = FileSelector(failed_txn_log_dir, recursive=False)
+    failed_txn_info_list = filesystem.get_file_info(failed_txn_file_selector)
+
+    for failed_txn_info in failed_txn_info_list:
+        try:
+            txn = Transaction.read(failed_txn_info.path, filesystem)
+            failed_txn_basename = posixpath.basename(failed_txn_info.path)
+            should_process = True
+            try:
+                if txn.state(catalog_root_normalized) == TransactionState.PURGED:
+                    should_process = False
+            except Exception:
+                logger.error("Could not check attribute")
+            if should_process:
+                # Process if the file is marked as currently cleaning.
+                txnid = txn.id
+
+                if txn.state(catalog_root_normalized) == TransactionState.FAILED:
+
+                    txnid = txn.id
+
+                    operations = txn["operations"]
+                    known_write_paths = chain.from_iterable(
+                        (op["metafile_write_paths"] + op["locator_write_paths"])
+                        for op in operations
+                    )
+
+                    for write_path in known_write_paths:
+                        full_path = posixpath.join(catalog_root_normalized, write_path)
+                        try:
+                            filesystem.delete_file(full_path)
+                        except Exception as e:
+                            logger.error(f"Failed to delete file '{full_path}': {e}")
+
+                    new_filename = f"{txnid}"
+
+                    new_failed_txn_log_file_path = posixpath.join(
+                        failed_txn_log_dir, new_filename
+                    )
+                    running_txn_log_path = posixpath.join(
+                        running_txn_log_dir, new_filename
+                    )
+
+                    os.delete(running_txn_log_path)
+
+                    os.rename(failed_txn_info.path, new_failed_txn_log_file_path)
+                    logger.debug(
+                        f"Cleaned up failed transaction: {failed_txn_basename}"
+                    )
+
+        except Exception as e:
+            logger.error(
+                f"Could not read transaction '{failed_txn_info.path}', skipping: {e}"
+            )
+
+
+def janitor_job(catalog_root_dir: str) -> None:
+    janitor_delete_timed_out_transaction(catalog_root_dir)
+    janitor_remove_files_in_failed(catalog_root_dir)