tab-cli 0.1.1__py3-none-any.whl

tab_cli/handlers/base.py

@@ -0,0 +1,259 @@
+"""Base classes for table reading and writing."""
+
+import os
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+import polars as pl
+from rich import box
+from rich.progress import Progress
+from rich.table import Table
+
+from tab_cli.formats.base import FormatHandler
+from tab_cli.storage.base import StorageBackend
+from tab_cli.style import _ALT_ROW_STYLE_0, _ALT_ROW_STYLE_1, _KEY_STYLE, _VAL_STYLE
+
+
+@dataclass
+class TableSchema:
+    """Schema information for a table."""
+
+    columns: list[tuple[str, pl.DataType]]
+
+    def __rich__(self) -> Table:
+        table = Table(
+            show_header=False,
+            box=box.SIMPLE_HEAD,
+            row_styles=[_ALT_ROW_STYLE_0, _ALT_ROW_STYLE_1],
+        )
+        table.add_column(style=_KEY_STYLE)
+        table.add_column(style=_VAL_STYLE)
+        for name, dtype in self.columns:
+            table.add_row(name, str(dtype))
+        return table
+
+
+@dataclass
+class TableSummary:
+    """Summary information for a table."""
+
+    file_size: int
+    num_rows: int
+    num_columns: int
+    extra: dict[str, str | int | float] | None = None
+
+    def __rich__(self) -> Table:
+        def format_size(size: int) -> str:
+            s: float = size
+            for unit in ["B", "KiB", "MiB", "GiB", "TiB"]:
+                if s < 1024:
+                    return f"{s:.1f} {unit}" if unit != "B" else f"{int(s)} {unit}"
+                s /= 1024
+            return f"{s:.1f} PiB"
+
+        table = Table(
+            show_header=False,
+            box=box.SIMPLE_HEAD,
+            row_styles=["", _ALT_ROW_STYLE],
+        )
+        table.add_column(style=_KEY_STYLE)
+        table.add_column(style=_VAL_STYLE)
+
+        table.add_row("File size", format_size(self.file_size))
+        table.add_row("Rows", f"{self.num_rows:,}")
+        table.add_row("Columns", str(self.num_columns))
+
+        if self.extra:
+            for key, value in self.extra.items():
+                table.add_row(key, str(value))
+
+        return table
+
+
+class TableReader:
+    """Reads tabular data by composing a StorageBackend and FormatHandler."""
+
+    def __init__(self, backend: StorageBackend, format: FormatHandler):
+        self.backend = backend
+        self.format = format
+
+    def read(self, url: str, limit: int | None = None, offset: int = 0) -> pl.LazyFrame:
+        if self.backend.is_directory(url):
+            lf = self._read_directory(url)
+        else:
+            polars_uri = self.backend.normalize_for_polars(url)
+            storage_options = self.backend.storage_options(url)
+            lf = self.format.scan(polars_uri, storage_options=storage_options)
+
+        if offset > 0:
+            lf = lf.slice(offset, length=limit)
+        elif limit is not None:
+            lf = lf.head(limit)
+        return lf
+
+    def _read_directory(self, url: str) -> pl.LazyFrame:
+        """Read all files in a directory."""
+        extension = self.format.extension()
+        storage_options = self.backend.storage_options(url)
+
+        if self.format.supports_glob():
+            # Use native glob support
+            polars_uri = self.backend.normalize_for_polars(url)
+            glob_pattern = os.path.join(polars_uri, "**", f"*{extension}")
+            return self.format.scan(glob_pattern, storage_options=storage_options)
+        else:
+            # Manual concatenation for formats without glob support
+            files = list(self.backend.list_files(url, extension))
+            if not files:
+                raise ValueError(f"No {extension} files found in {url}")
+            frames = [
+                self.format.scan(
+                    self.backend.normalize_for_polars(f.url),
+                    storage_options=self.backend.storage_options(f.url),
+                )
+                for f in files
+            ]
+            return pl.concat(frames, how="vertical")
+
+    def schema(self, url: str) -> TableSchema:
+        if self.backend.is_directory(url):
+            # Get schema from first file
+            files = list(self.backend.list_files(url, self.format.extension()))
+            if not files:
+                raise ValueError(f"No {self.format.extension()} files found in {url}")
+            url = files[0].url
+        polars_uri = self.backend.normalize_for_polars(url)
+        storage_options = self.backend.storage_options(url)
+        columns = self.format.collect_schema(polars_uri, storage_options=storage_options)
+        return TableSchema(columns=columns)
+
+    def summary(self, url: str) -> TableSummary:
+        if self.backend.is_directory(url):
+            return self._summary_directory(url)
+        else:
+            return self._summary_single(url)
+
+    def _summary_single(self, url: str) -> TableSummary:
+        file_size = self.backend.size(url)
+        polars_uri = self.backend.normalize_for_polars(url)
+        storage_options = self.backend.storage_options(url)
+        num_rows = self.format.count_rows(polars_uri, storage_options=storage_options)
+        schema = self.format.collect_schema(polars_uri, storage_options=storage_options)
+        num_columns = len(schema)
+        extra = self.format.extra_summary(url)
+        return TableSummary(
+            file_size=file_size,
+            num_rows=num_rows,
+            num_columns=num_columns,
+            extra=extra,
+        )
+
+    def _summary_directory(self, url: str) -> TableSummary:
+        """Aggregate summary from all files in directory."""
+        files = list(self.backend.list_files(url, self.format.extension()))
+        if not files:
+            raise ValueError(f"No {self.format.extension()} files found in {url}")
+
+        file_size = 0
+        num_rows = 0
+        num_columns: int | None = None
+
+        extra_numeric: dict[str, float] = {}
+        extra_strings: dict[str, set[str]] = {}
+
+        for file_info in files:
+            file_size += file_info.size
+            polars_uri = self.backend.normalize_for_polars(file_info.url)
+            storage_options = self.backend.storage_options(file_info.url)
+            num_rows += self.format.count_rows(polars_uri, storage_options=storage_options)
+
+            schema = self.format.collect_schema(polars_uri, storage_options=storage_options)
+            if num_columns is None:
+                num_columns = len(schema)
+            elif len(schema) != num_columns:
+                raise ValueError(f"Inconsistent column counts in {url}")
+
+            extra = self.format.extra_summary(file_info.url)
+            if extra:
+                for key, value in extra.items():
+                    if isinstance(value, (int, float)):
+                        extra_numeric[key] = extra_numeric.get(key, 0) + value
+                    else:
+                        extra_strings.setdefault(key, set()).add(str(value))
+
+        result_extra: dict[str, str | int | float] = {"Partitions": len(files)}
+        for key, value in extra_numeric.items():
+            if float(value).is_integer():
+                result_extra[key] = int(value)
+            else:
+                result_extra[key] = value
+
+        for key, values in extra_strings.items():
+            if len(values) == 1:
+                result_extra[key] = next(iter(values))
+            else:
+                result_extra[key] = ", ".join(sorted(values))
+
+        return TableSummary(
+            file_size=file_size,
+            num_rows=num_rows,
+            num_columns=num_columns or 0,
+            extra=result_extra,
+        )
+
+
+class TableWriter(ABC):
+    """Base class for writing tabular data."""
+
+    @abstractmethod
+    def extension(self) -> str:
+        """Return the file extension for this format."""
+        pass
+
+    @abstractmethod
+    def write(self, lf: pl.LazyFrame) -> Iterable[bytes]:
+        """Write LazyFrame to bytes (for streaming output)."""
+        pass
+
+    @abstractmethod
+    def write_to_single_file(self, lf: pl.LazyFrame, path: str) -> None:
+        """Write LazyFrame to a single file."""
+        pass
+
+    def write_to_path(self, lf: pl.LazyFrame, path: str, partitions: int | None = None) -> None:
+        """Write LazyFrame to a file or partitioned directory."""
+        if partitions is None:
+            with Progress() as progress:
+                task = progress.add_task("Writing...", total=1)
+                self.write_to_single_file(lf, path)
+                progress.update(task, completed=1)
+        else:
+            os.makedirs(path, exist_ok=True)
+            row_count = lf.select(pl.len()).collect().item()
+            rows_per_part = (row_count + partitions - 1) // partitions
+            with Progress() as progress:
+                task = progress.add_task("Writing partitions...", total=partitions)
+                for i in range(partitions):
+                    offset = i * rows_per_part
+                    if offset < row_count:
+                        part_lf = lf.slice(offset, rows_per_part)
+                        part_path = os.path.join(path, f"part-{i:05d}{self.extension()}")
+                        self.write_to_single_file(part_lf, part_path)
+                    progress.update(task, advance=1)
+
+
+class FormatWriter(TableWriter):
+    """TableWriter adapter for FormatHandler."""
+
+    def __init__(self, format: FormatHandler):
+        self._format = format
+
+    def extension(self) -> str:
+        return self._format.extension()
+
+    def write(self, lf: pl.LazyFrame) -> Iterable[bytes]:
+        return self._format.write(lf)
+
+    def write_to_single_file(self, lf: pl.LazyFrame, path: str) -> None:
+        self._format.write_to_single_file(lf, path)

tab_cli/handlers/cli_table.py

@@ -0,0 +1,55 @@
+import sys
+from collections.abc import Iterable
+
+from rich.table import Table
+from rich import box
+from rich.console import Console
+import polars as pl
+
+from tab_cli.handlers.base import TableWriter
+from tab_cli.style import _ALT_ROW_STYLE_0, _ALT_ROW_STYLE_1, _KEY_STYLE
+
+
+class CliTableFormatter(TableWriter):
+    def __init__(self, truncated: bool = False, svg_capture: bool = False):
+        self.truncated = truncated
+        self.svg_capture = svg_capture
+
+    def extension(self) -> str:
+        return ".txt"
+
+    def write(self, lf: pl.LazyFrame) -> Iterable[bytes]:
+
+        table = Table(
+            show_header=True,
+            header_style=_KEY_STYLE,
+            box=box.SIMPLE_HEAD,
+            row_styles=[_ALT_ROW_STYLE_0, _ALT_ROW_STYLE_1],
+        )
+
+        for col in lf.collect_schema().names():
+            table.add_column(col)
+
+        for batch in lf.collect_batches():
+            for row in batch.iter_rows():
+                table.add_row(*[str(v) if v is not None else "" for v in row])
+
+        if self.truncated:
+            table.add_row(*["..." for _ in lf.collect_schema().names()])
+
+        if self.svg_capture:
+            console = Console(record=True, width=80)
+            console.print(table)
+            svg = console.export_svg()
+            print(svg, file=sys.stderr)
+        else:
+            console = Console()
+            with console.capture() as capture:
+                console.print(table)
+            yield capture.get().encode("utf-8")
+
+    def write_to_single_file(self, lf: pl.LazyFrame, path: str) -> None:
+        """Write a LazyFrame to a single text file."""
+        with open(path, "wb") as f:
+            for chunk in self.write(lf):
+                f.write(chunk)

tab_cli/storage/__init__.py

@@ -0,0 +1,83 @@
+"""Storage backends for filesystem abstraction."""
+
+import os
+from dataclasses import dataclass
+from loguru import logger
+from urllib.parse import urlparse
+
+from tab_cli.storage.base import FileInfo, StorageBackend
+from tab_cli.storage.local import LocalBackend
+from tab_cli.url_parser import parse_url, ParsedUrl
+
+__all__ = [
+    "FileInfo",
+    "StorageBackend",
+    "LocalBackend",
+    "get_backend",
+]
+
+
+def get_backend(url: str) -> StorageBackend:
+    """Get the appropriate storage backend for a URL.
+
+    Supports:
+    - Local paths (no scheme or file://)
+    - s3:// - AWS S3 (requires s3fs)
+    - gs:// - Google Cloud Storage (requires gcsfs)
+    - az:// - Azure Blob Storage (requires adlfs)
+    - abfs://, abfss:// - Azure Data Lake Storage Gen2 (requires adlfs)
+    - Any other fsspec-supported protocol
+
+    For az:// URLs, the interpretation of the URL authority depends on the
+    --az-url-authority-is-account global flag:
+    - If set: authority is the storage account name
+      - az://account/container/path
+      - az:///container/path (account from AZURE_STORAGE_ACCOUNT)
+    - If not set (default): authority is the container name
+      - az://container/path (standard adlfs behavior)
+    """
+    parsed = parse_url(url)
+    logger.debug(f"Accessing data from\n"
+                 f" - Protocol: [bold]{parsed.scheme}[/]\n"
+                 f" - Account: {parsed.account}\n"
+                 f" - Bucket: {parsed.bucket}\n"
+                 f" - Path: {parsed.path}"
+                 )
+
+    # Local filesystem
+    if parsed.scheme == "file" or not parsed.scheme:
+        return LocalBackend()
+
+    elif parsed.scheme == "az":
+        from tab_cli import config
+        from tab_cli.storage.az import AzBackend
+
+        return AzBackend(
+            account=parsed.account,
+            container=parsed.bucket,
+            az_url_authority_is_account=config.config.az_url_authority_is_account,
+        )
+
+    elif parsed.scheme in {"abfs", "abfss"}:
+        from tab_cli.storage.az import AzBackend
+
+        # abfs/abfss always uses account in URL or env
+        return AzBackend(
+            account=parsed.account,
+            container=parsed.bucket,
+        )
+
+    # Google Cloud Storage
+    elif parsed.scheme == "gs":
+        from tab_cli.storage.gcloud import GcloudBackend
+        return GcloudBackend()
+
+    # AWS S3
+    elif parsed.scheme == "s3":
+        from tab_cli.storage.aws import AwsBackend
+        return AwsBackend()
+
+    # All other protocols via fsspec
+    from tab_cli.storage.fsspec import FsspecBackend
+
+    return FsspecBackend(parsed.scheme)

tab_cli/storage/aws.py

@@ -0,0 +1,223 @@
+import os
+from enum import Enum
+
+from typing import BinaryIO, Iterator
+from loguru import logger
+
+from tab_cli.storage.base import StorageBackend, FileInfo
+from tab_cli.url_parser import parse_url
+
+
+class AwsAuthMethod(Enum):
+    EXPLICIT_KEYS = 1    # AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY
+    PROFILE = 2          # AWS_PROFILE or default profile (handles SSO, assume role, etc.)
+    ANONYMOUS = 3        # Public buckets
+
+
+class AwsBackend(StorageBackend):
+    """Storage backend for AWS S3.
+
+    URL format: s3://bucket/path
+
+    Authentication is handled by boto3's credential chain via s3fs:
+    1. Explicit keys from environment (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN)
+    2. Profile-based auth (AWS_PROFILE or default) - handles:
+       - Shared credentials file (~/.aws/credentials)
+       - AWS config file (~/.aws/config)
+       - SSO credentials (from `aws sso login`)
+       - Assume role
+       - Container credentials (ECS/EKS)
+       - EC2 instance metadata
+    3. Anonymous access (for public buckets, if requested)
+    """
+
+    def __init__(self, anon: bool = False) -> None:
+        """Initialize the AWS S3 storage backend.
+
+        Args:
+            anon: If True, use anonymous access (for public buckets only).
+        """
+        try:
+            import s3fs
+        except ImportError as e:
+            raise ImportError("Package 's3fs' is required for s3:// URLs. Install with: pip install s3fs") from e
+
+        self.s3fs = s3fs
+        self.fs = None
+        self.anon = anon
+
+        # Get profile from environment
+        self.profile = os.environ.get("AWS_PROFILE")
+        self.region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
+
+        # Check for explicit credentials in environment
+        self.access_key = os.environ.get("AWS_ACCESS_KEY_ID")
+        self.secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
+        self.session_token = os.environ.get("AWS_SESSION_TOKEN")
+
+        if anon:
+            # Anonymous access for public buckets
+            logger.debug("Using anonymous access for S3")
+            self.fs = self.s3fs.S3FileSystem(anon=True)
+            self.method = AwsAuthMethod.ANONYMOUS
+            return
+
+        # 1. Try explicit credentials from environment
+        if self.access_key and self.secret_key:
+            logger.debug("Authenticating to S3 using explicit credentials from environment")
+            try:
+                self.fs = self.s3fs.S3FileSystem(
+                    key=self.access_key,
+                    secret=self.secret_key,
+                    token=self.session_token,
+                )
+                self.method = AwsAuthMethod.EXPLICIT_KEYS
+                return
+            except Exception as e:
+                logger.debug("Explicit credentials authentication failed: {}", e)
+
+        # 2. Fall back to profile-based auth (boto3 credential chain)
+        # This handles: ~/.aws/credentials, ~/.aws/config, SSO, assume role, instance metadata
+        profile_desc = f"profile '{self.profile}'" if self.profile else "default credential chain"
+        logger.debug("Authenticating to S3 using {}", profile_desc)
+        try:
+            self.fs = self.s3fs.S3FileSystem(profile=self.profile)
+            self.method = AwsAuthMethod.PROFILE
+            return
+        except Exception as e:
+            logger.debug("Profile-based authentication failed: {}", e)
+
+        if self.fs is None:
+            raise ValueError(
+                "Could not authenticate to AWS S3. "
+                "Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY, "
+                "configure ~/.aws/credentials, run 'aws configure', "
+                "or run 'aws sso login'."
+            )
+
+    def normalize_for_polars(self, url: str) -> str:
+        """Normalize URL to a format Polars understands.
+
+        Polars expects s3://bucket/path format.
+
+        Returns:
+            Normalized URL in s3://bucket/path format.
+        """
+        parsed = parse_url(url)
+        return f"s3://{parsed.bucket}/{parsed.path}"
+
+    def storage_options(self, url: str) -> dict[str, str] | None:
+        """Return storage options for Polars S3 access.
+
+        Returns:
+            Dict with appropriate authentication options for S3.
+            Includes both s3fs-style keys and Rust object_store keys for compatibility.
+        """
+        if self.method == AwsAuthMethod.ANONYMOUS:
+            return {
+                "anon": True,
+                "aws_skip_signature": "true",
+            }
+
+        if self.method == AwsAuthMethod.EXPLICIT_KEYS:
+            opts = {
+                # s3fs keys
+                "key": self.access_key,
+                "secret": self.secret_key,
+                # object_store keys
+                "aws_access_key_id": self.access_key,
+                "aws_secret_access_key": self.secret_key,
+            }
+            if self.session_token:
+                opts["token"] = self.session_token
+                opts["aws_session_token"] = self.session_token
+            if self.region:
+                opts["client_kwargs"] = {"region_name": self.region}
+                opts["aws_region"] = self.region
+            return opts
+
+        if self.method == AwsAuthMethod.PROFILE:
+            # For profile-based auth, we need to fetch the resolved credentials
+            # so Polars can use them (Polars doesn't understand profiles directly)
+            try:
+                credentials = self._get_credentials_from_session()
+                if credentials:
+                    opts = {
+                        # s3fs keys
+                        "key": credentials["access_key"],
+                        "secret": credentials["secret_key"],
+                        # object_store keys
+                        "aws_access_key_id": credentials["access_key"],
+                        "aws_secret_access_key": credentials["secret_key"],
+                    }
+                    if credentials.get("token"):
+                        opts["token"] = credentials["token"]
+                        opts["aws_session_token"] = credentials["token"]
+                    if self.region:
+                        opts["client_kwargs"] = {"region_name": self.region}
+                        opts["aws_region"] = self.region
+                    return opts
+            except Exception as e:
+                logger.debug("Failed to resolve credentials from session: {}", e)
+
+            # Fallback: just pass the profile and hope Polars/fsspec can handle it
+            opts = {}
+            if self.profile:
+                opts["profile"] = self.profile
+            if self.region:
+                opts["client_kwargs"] = {"region_name": self.region}
+                opts["aws_region"] = self.region
+            return opts if opts else None
+
+        return None
+
+    def _get_credentials_from_session(self) -> dict | None:
+        """Get resolved credentials from boto3 session."""
+        try:
+            import boto3
+            session = boto3.Session(profile_name=self.profile)
+            credentials = session.get_credentials()
+            if credentials:
+                frozen = credentials.get_frozen_credentials()
+                return {
+                    "access_key": frozen.access_key,
+                    "secret_key": frozen.secret_key,
+                    "token": frozen.token,
+                }
+        except Exception:
+            pass
+        return None
+
+    def _to_internal(self, url: str) -> str:
+        """Convert URL to internal path for s3fs operations."""
+        parsed = parse_url(url)
+        return f"{parsed.bucket}/{parsed.path}"
+
+    def _to_uri(self, internal_path: str) -> str:
+        """Convert internal path back to s3:// URL."""
+        return f"s3://{internal_path}"
+
+    def open(self, url: str) -> BinaryIO:
+        return self.fs.open(self._to_internal(url), "rb")
+
+    def list_files(self, url: str, extension: str) -> Iterator[FileInfo]:
+        internal_path = self._to_internal(url)
+        pattern = f"{internal_path}/**/*{extension}"
+        for path in sorted(self.fs.glob(pattern)):
+            info = self.fs.info(path)
+            yield FileInfo(url=self._to_uri(path), size=info["size"])
+
+    def size(self, url: str) -> int:
+        return self.fs.size(self._to_internal(url))
+
+    def is_directory(self, url: str) -> bool:
+        path = self._to_internal(url)
+        try:
+            info = self.fs.info(path)
+            return info.get("type") == "directory"
+        except FileNotFoundError:
+            try:
+                contents = self.fs.ls(path, detail=False)
+                return len(contents) > 0
+            except Exception:
+                return False