arrows 0.2.0__py3-none-any.whl

arrows/__init__.py

@@ -0,0 +1,159 @@
+"""arrows — an Arrow-native ETL toolkit with pluggable components.
+
+Nothing heavy is imported at module load. Components (and the third-party
+libraries they need) are pulled in the first time they are used::
+
+    import arrows
+
+    arrows.load('s3', 'redshift')          # load exactly what this job needs
+    arrow = arrows.redshift.fetch_arrow('select 1')
+
+``arrows.list_components()`` shows what is available, including components
+contributed by other packages through the ``arrows.components`` entry point.
+"""
+
+from __future__ import annotations
+
+import importlib
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _metadata_version
+from typing import TYPE_CHECKING, Any
+
+# Registering the built-in specs costs nothing: only metadata is imported here.
+from . import components as _components  # noqa: F401
+from .core.component import Component, ComponentSpec, HealthStatus
+from .core.errors import (
+    ArrowsError,
+    ComponentNotFoundError,
+    ComponentNotLoadedError,
+    ComponentSetupError,
+    MissingDependencyError,
+    MissingSecretError,
+)
+from .core.registry import ENTRY_POINT_GROUP, register, specs
+from .core.secrets import PromptProvider, Secret, SecretStore
+from .core.session import (
+    Session,
+    close,
+    configure,
+    default_session,
+    get,
+    health,
+    is_loaded,
+    load,
+    login,
+    unload,
+    use_session,
+)
+
+#: pyproject.toml is the single source of truth; read it back from the installed
+#: metadata rather than keeping a second copy here that can drift from the tag.
+try:
+    __version__ = _metadata_version('arrows')
+except PackageNotFoundError:  # running from a source tree that was never installed
+    __version__ = '0.0.0.dev0'
+
+#: Submodules exposed as attributes but imported on first access.
+_LAZY_MODULES = {
+    'auth',
+    'gmail',
+    'google_sheets',
+    'redshift',
+    's3',
+    'sqlite',
+    'template_renderer',
+    'utils',
+}
+
+if TYPE_CHECKING:  # give editors and type checkers the real modules
+    from . import auth, gmail, google_sheets, redshift, s3, sqlite, template_renderer, utils
+
+__all__ = [
+    'ENTRY_POINT_GROUP',
+    'ArrowsError',
+    'Component',
+    'ComponentNotFoundError',
+    'ComponentNotLoadedError',
+    'ComponentSetupError',
+    'ComponentSpec',
+    'HealthStatus',
+    'MissingDependencyError',
+    'MissingSecretError',
+    'PromptProvider',
+    'Secret',
+    'SecretStore',
+    'Session',
+    'auth',
+    'close',
+    'configure',
+    'default_session',
+    'get',
+    'gmail',
+    'google_sheets',
+    'health',
+    'is_loaded',
+    'list_components',
+    'load',
+    'load_credentials',
+    'login',
+    'redshift',
+    'register',
+    's3',
+    'sqlite',
+    'template_renderer',
+    'unload',
+    'use_session',
+    'utils',
+]
+
+
+def __getattr__(name: str) -> Any:
+    """PEP 562 lazy submodule import.
+
+    ``arrows.redshift`` works without ``import arrows`` dragging in psycopg2,
+    boto3, awswrangler and the Google client for every process that only wanted
+    to send an email.
+    """
+    if name in _LAZY_MODULES:
+        module = importlib.import_module(f'.{name}', __name__)
+        globals()[name] = module
+        return module
+    raise AttributeError(f'module {__name__!r} has no attribute {name!r}')
+
+
+def __dir__() -> list[str]:
+    return sorted(set(__all__) | set(globals()))
+
+
+def list_components() -> list[dict[str, Any]]:
+    """Describe every registered component: name, dependencies, load state."""
+    session = default_session()
+    return [
+        {
+            'name': spec.name,
+            'summary': spec.summary,
+            'depends_on': list(spec.depends_on),
+            'install_hint': spec.install_hint,
+            'loaded': session.is_loaded(spec.name),
+        }
+        for spec in sorted(specs().values(), key=lambda s: s.name)
+    ]
+
+
+def load_credentials(*names: str) -> Session:
+    """Deprecated alias for :func:`load`.
+
+    The old behaviour — loading AWS, Redshift and Google unconditionally — is
+    kept only when called with no arguments, and warns.
+    """
+    import warnings
+
+    if not names:
+        warnings.warn(
+            'load_credentials() without arguments loads every component and fails if any '
+            "credential is missing. Prefer arrows.load('s3', 'redshift').",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        names = ('s3', 'redshift', 'google')
+    return load(*names)

arrows/__main__.py

@@ -0,0 +1,6 @@
+import sys
+
+from .cli import main
+
+if __name__ == '__main__':
+    sys.exit(main())

arrows/auth.py

@@ -0,0 +1,39 @@
+"""Backwards-compatible credential loading.
+
+Superseded by :mod:`arrows.core.secrets` and :func:`arrows.load`. These
+functions remain so existing scripts keep running; each one loads the matching
+component and warns once.
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+from .core.session import default_session
+
+__all__ = ['load_aws_credentials', 'load_google_credentials', 'load_redshift_credentials']
+
+
+def _deprecated(old: str, new: str) -> None:
+    warnings.warn(f'arrows.auth.{old}() is deprecated; use {new} instead.', DeprecationWarning, stacklevel=3)
+
+
+def load_aws_credentials() -> Any:
+    _deprecated('load_aws_credentials', "arrows.load('s3')")
+    return default_session().get('s3')
+
+
+def load_redshift_credentials() -> Any:
+    _deprecated('load_redshift_credentials', "arrows.load('redshift')")
+    return default_session().get('redshift')
+
+
+def load_google_credentials() -> Any:
+    _deprecated('load_google_credentials', "arrows.load('google')")
+    return default_session().get('google')
+
+
+def _get_google_credentials() -> Any:
+    """Legacy accessor for the raw google-auth credentials object."""
+    return default_session().get('google').credentials

arrows/cli.py

@@ -0,0 +1,102 @@
+"""``arrows`` command line: see what exists, and whether it works.
+
+arrows components            list every registered component
+arrows doctor s3 redshift    load them and run their health checks
+arrows secrets redshift      show which required secrets resolve (redacted)
+arrows login redshift --save type the missing ones, store them in the keychain
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from . import list_components
+from .core import registry
+from .core.session import default_session
+
+
+def _cmd_components(args: argparse.Namespace) -> int:
+    rows = list_components()
+    width = max((len(row['name']) for row in rows), default=4)
+    for row in rows:
+        marker = '*' if row['loaded'] else ' '
+        depends = f' (needs {", ".join(row["depends_on"])})' if row['depends_on'] else ''
+        hint = f'  [{row["install_hint"]}]' if row['install_hint'] else ''
+        print(f'{marker} {row["name"]:<{width}}  {row["summary"]}{depends}{hint}')
+    print('\n* = loaded in this session')
+    return 0
+
+
+def _cmd_doctor(args: argparse.Namespace) -> int:
+    session = default_session()
+    names = args.components or registry.names()
+    exit_code = 0
+    for name in names:
+        try:
+            session.load(name)
+            status = session.get(name).health_check()
+            print(f'{"ok  " if status.ok else "FAIL"}  {name:<14} {status.detail}')
+            exit_code |= 0 if status.ok else 1
+        except Exception as exc:
+            print(f'FAIL  {name:<14} {type(exc).__name__}: {exc}')
+            exit_code = 1
+    return exit_code
+
+
+def _cmd_secrets(args: argparse.Namespace) -> int:
+    session = default_session()
+    print(f'providers: {", ".join(p.name for p in session.secrets.providers)}\n')
+    exit_code = 0
+    for name in args.components or registry.names():
+        spec = registry.get_spec(name)
+        component = spec.build()
+        print(f'{spec.name}:')
+        for key in component.requires:
+            found = session.secrets.get(key)
+            print(f'  {"ok  " if found else "MISS"}  {key:<32} {found.source if found else "-"}')
+            exit_code |= 0 if found else 1
+        for key in component.optional:
+            found = session.secrets.get(key)
+            print(f'  {"ok  " if found else "--  "}  {key:<32} {(found.source if found else "optional")}')
+        if not component.requires and not component.optional:
+            print('  (no secrets required)')
+    return exit_code
+
+
+def _cmd_login(args: argparse.Namespace) -> int:
+    session = default_session()
+    try:
+        session.login(*args.components, save=args.save, include_optional=args.all)
+    except Exception as exc:
+        print(f'{type(exc).__name__}: {exc}', file=sys.stderr)
+        return 1
+    return _cmd_doctor(args)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog='arrows', description=__doc__.splitlines()[0])
+    subparsers = parser.add_subparsers(dest='command', required=True)
+
+    subparsers.add_parser('components', help='list registered components').set_defaults(func=_cmd_components)
+
+    doctor = subparsers.add_parser('doctor', help='load components and run health checks')
+    doctor.add_argument('components', nargs='*', help='component names (default: all)')
+    doctor.set_defaults(func=_cmd_doctor)
+
+    secrets = subparsers.add_parser('secrets', help='report which secrets resolve, without revealing them')
+    secrets.add_argument('components', nargs='*', help='component names (default: all)')
+    secrets.set_defaults(func=_cmd_secrets)
+
+    login = subparsers.add_parser('login', help='prompt for missing secrets, then load and check')
+    login.add_argument('components', nargs='+', help='component names')
+    login.add_argument('--save', action='store_true', help='remember the answers in the OS keychain')
+    login.add_argument('--all', action='store_true', help='also ask for optional keys')
+    login.set_defaults(func=_cmd_login)
+
+    args = parser.parse_args(argv)
+    return args.func(args)
+
+
+if __name__ == '__main__':
+    sys.exit(main())

arrows/components/__init__.py

@@ -0,0 +1,69 @@
+"""Built-in components.
+
+Only the :class:`ComponentSpec` metadata is imported here; the modules that
+define the components — and their heavy third-party dependencies — are imported
+the first time a component is actually loaded.
+
+To add a component:
+
+1. Copy ``sqlite.py``, implement ``setup``/``close``/``health_check``.
+2. Declare ``requires`` (secret keys) and ``depends_on`` (other components).
+3. Export ``COMPONENT`` and a ``SPEC``.
+4. Register the ``SPEC`` below — or, from a separate distribution, advertise it
+   through the ``arrows.components`` entry-point group; nothing in arrows needs
+   to change.
+"""
+
+from __future__ import annotations
+
+from ..core.component import ComponentSpec
+from ..core.registry import register
+
+__all__ = ['BUILTIN_SPECS']
+
+BUILTIN_SPECS: tuple[ComponentSpec, ...] = (
+    ComponentSpec(
+        name='aws',
+        module='arrows.components.aws',
+        summary='boto3 session + DuckDB S3 credential chain (shared by s3 and redshift)',
+    ),
+    ComponentSpec(
+        name='s3',
+        module='arrows.components.s3',
+        summary='S3 datasets in Parquet, queryable with DuckDB/Polars',
+        depends_on=('aws',),
+    ),
+    ComponentSpec(
+        name='redshift',
+        module='arrows.components.redshift',
+        summary='Amazon Redshift: query, UNLOAD to S3, COPY from S3',
+        depends_on=('aws',),
+    ),
+    ComponentSpec(
+        name='google',
+        module='arrows.components.google',
+        summary='Google OAuth credentials shared by gmail and google_sheets',
+        aliases=('google_auth',),
+    ),
+    ComponentSpec(
+        name='google_sheets',
+        module='arrows.components.google_sheets',
+        summary='Read and write Google Sheets as Arrow/Polars/DuckDB',
+        depends_on=('google',),
+        aliases=('sheets', 'googlesheets'),
+    ),
+    ComponentSpec(
+        name='gmail',
+        module='arrows.components.gmail',
+        summary='Send email through the Gmail API',
+        depends_on=('google',),
+    ),
+    ComponentSpec(
+        name='sqlite',
+        module='arrows.components.sqlite',
+        summary='Local SQLite files, read and written through DuckDB (reference component)',
+    ),
+)
+
+for _spec in BUILTIN_SPECS:
+    register(_spec, replace=True)

arrows/components/aws.py

@@ -0,0 +1,164 @@
+"""AWS credentials, shared by the s3 and redshift components.
+
+Best practice applied here: arrows does **not** invent a credential format. It
+defers to boto3's own resolution chain (env vars, ``~/.aws/credentials``, SSO,
+``AWS_PROFILE``, EC2/ECS/EKS instance roles), which is the only path that
+supports short-lived, automatically rotated credentials. Explicit keys from the
+:class:`SecretStore` are honoured when present, for laptops that still use them.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from pathlib import Path
+from typing import Any, ClassVar
+
+from ..core.component import Component, HealthStatus
+from ..core.engine import get_duckdb, quote_literal
+from ..core.secrets import SecretStore
+
+EXPLICIT_KEYS = ('AWS_ACCESS_KEY_ID', 'AWS_SECRET_ACCESS_KEY', 'AWS_SESSION_TOKEN')
+
+
+class LegacyAwsFileProvider:
+    """Reads the historical ``~/.credentials/aws_credentials.txt`` layout.
+
+    Kept so existing setups keep working; prefer a real ``~/.aws/credentials``
+    profile or SSO, which boto3 refreshes on its own.
+    """
+
+    def __init__(self, path: str | Path):
+        self.path = Path(path).expanduser()
+        self.name = f'legacy-aws-file({self.path})'
+        self._cache: dict[str, str] | None = None
+
+    def _load(self) -> dict[str, str]:
+        if self._cache is not None:
+            return self._cache
+        values: dict[str, str] = {}
+        if self.path.is_file():
+            from ..core.secrets import _warn_if_world_readable
+
+            _warn_if_world_readable(self.path)
+            for line in self.path.read_text(encoding='utf-8').splitlines():
+                line = line.strip()
+                if not line or line.startswith('[') or '=' not in line:
+                    continue
+                key, _, value = line.partition('=')
+                key = key.strip().upper()
+                if key in EXPLICIT_KEYS:
+                    values[key] = value.strip()
+        self._cache = values
+        return values
+
+    def get(self, key: str) -> str | None:
+        return self._load().get(key)
+
+
+class AwsComponent(Component):
+    """Owns the boto3 session and registers DuckDB's S3 secret."""
+
+    name = 'aws'
+    optional = (*EXPLICIT_KEYS, 'AWS_REGION', 'AWS_PROFILE')
+    login_args: ClassVar[dict[str, str]] = {
+        'access_key_id': 'AWS_ACCESS_KEY_ID',
+        'secret_access_key': 'AWS_SECRET_ACCESS_KEY',
+        'session_token': 'AWS_SESSION_TOKEN',
+        'region': 'AWS_REGION',
+        'profile': 'AWS_PROFILE',
+    }
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.boto3_session: Any = None
+        self.region: str | None = None
+        self._explicit_credentials = False
+
+    def setup(self, secrets: SecretStore) -> None:
+        boto3 = self.import_module('boto3')
+
+        # Backwards compatibility only, and only when the file is really there:
+        # a production store built without file providers must stay that way.
+        credentials_dir = (
+            Path(self.session.config.credentials_dir).expanduser()
+            if self.session
+            else Path('~/.credentials').expanduser()
+        )
+        legacy_file = credentials_dir / 'aws_credentials.txt'
+        if legacy_file.is_file():
+            secrets.add_provider(LegacyAwsFileProvider(legacy_file))
+
+        self.region = self.secret('AWS_REGION') or self.secret('AWS_DEFAULT_REGION')
+        access_key = self.secret('AWS_ACCESS_KEY_ID')
+        kwargs: dict[str, Any] = {'region_name': self.region} if self.region else {}
+        self._explicit_credentials = bool(access_key)
+        if access_key:
+            kwargs.update(
+                aws_access_key_id=access_key,
+                aws_secret_access_key=self.secret('AWS_SECRET_ACCESS_KEY'),
+                aws_session_token=self.secret('AWS_SESSION_TOKEN'),
+            )
+        elif profile := self.secret('AWS_PROFILE'):
+            kwargs['profile_name'] = profile
+
+        self.boto3_session = boto3.Session(**kwargs)
+        self._register_duckdb_secret()
+
+    def _register_duckdb_secret(self) -> None:
+        """Give DuckDB the same identity boto3 resolved.
+
+        Two paths, and the difference matters:
+
+        * When boto3 resolved credentials from its own chain (profile, SSO,
+          instance role), DuckDB is pointed at that same chain — no key material
+          is written into SQL, and DuckDB refreshes on its own.
+        * When the credentials came from the arrows SecretStore, DuckDB cannot
+          see that store, so the keys have to be handed over as a literal
+          secret. They are then visible in ``duckdb_secrets()``, which is the
+          cost of that setup and a reason to prefer a real AWS profile.
+
+        A failure here is not fatal: boto3 and pyarrow still work.
+        """
+        import warnings
+
+        credentials = self.boto3_session.get_credentials()
+        if credentials is None:
+            warnings.warn('No AWS credentials resolved; DuckDB S3 access is unconfigured.', stacklevel=2)
+            return
+
+        region = f', REGION {quote_literal(self.region)}' if self.region else ''
+        if self._explicit_credentials:
+            frozen = credentials.get_frozen_credentials()
+            token = f', SESSION_TOKEN {quote_literal(frozen.token)}' if frozen.token else ''
+            body = (
+                f'TYPE s3, KEY_ID {quote_literal(frozen.access_key)}, '
+                f'SECRET {quote_literal(frozen.secret_key)}{token}{region}'
+            )
+        else:
+            body = f'TYPE s3, PROVIDER credential_chain{region}'
+
+        try:
+            get_duckdb().execute(f'CREATE OR REPLACE SECRET arrows_s3 ({body});')
+        except Exception as exc:
+            warnings.warn(f'Could not register the DuckDB S3 secret: {exc}', stacklevel=2)
+
+    def client(self, service: str, **kwargs) -> Any:
+        return self.boto3_session.client(service, **kwargs)
+
+    def health_check(self) -> HealthStatus:
+        if self.boto3_session is None:
+            return HealthStatus(self.name, False, 'no boto3 session')
+        try:
+            identity = self.boto3_session.client('sts').get_caller_identity()
+            return HealthStatus(self.name, True, f'arn={identity.get("Arn", "?")}')
+        except Exception as exc:
+            return HealthStatus(self.name, False, f'{type(exc).__name__}: {exc}')
+
+    def close(self) -> None:
+        # Teardown is best effort: the connection may already be gone.
+        with contextlib.suppress(Exception):
+            get_duckdb().execute('DROP SECRET IF EXISTS arrows_s3;')
+        self.boto3_session = None
+
+
+COMPONENT = AwsComponent

arrows/components/gmail.py

@@ -0,0 +1,33 @@
+"""Gmail sending, on top of the shared google component."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ..core.component import Component, HealthStatus
+from ..core.secrets import SecretStore
+
+
+class GmailComponent(Component):
+    name = 'gmail'
+    depends_on = ('google',)
+
+    def setup(self, secrets: SecretStore) -> None:
+        self.import_module('googleapiclient.discovery')
+
+    @property
+    def service(self) -> Any:
+        return self.session.get('google').service('gmail', 'v1')
+
+    def address(self) -> str:
+        """The authenticated mailbox address."""
+        return self.service.users().getProfile(userId='me').execute().get('emailAddress', '')
+
+    def health_check(self) -> HealthStatus:
+        try:
+            return HealthStatus(self.name, True, f'mailbox={self.address()}')
+        except Exception as exc:
+            return HealthStatus(self.name, False, f'{type(exc).__name__}: {exc}')
+
+
+COMPONENT = GmailComponent

arrows/components/google.py

@@ -0,0 +1,115 @@
+"""Google OAuth credentials, shared by the gmail and google_sheets components.
+
+Best practice applied here: the token is refreshed **lazily and in memory**.
+The previous implementation refreshed at import time (a network call as a side
+effect of ``import arrows``) and then copied the whole token document into
+``os.environ``, where every subprocess inherited it. Here the credential object
+stays in the component, and the short-lived access token is re-issued only when
+it actually expires.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+from pathlib import Path
+from typing import Any, ClassVar
+
+from ..core.component import Component, HealthStatus
+from ..core.engine import get_duckdb, quote_literal
+from ..core.secrets import SecretStore
+
+DEFAULT_SCOPES = (
+    'https://www.googleapis.com/auth/spreadsheets',
+    'https://www.googleapis.com/auth/drive',
+    'https://www.googleapis.com/auth/gmail.send',
+)
+
+
+class GoogleComponent(Component):
+    """Holds refreshable Google OAuth credentials and builds API services."""
+
+    name = 'google'
+    requires = ('GOOGLE_TOKEN_JSON',)
+    optional = ('GOOGLE_SCOPES',)
+    login_args: ClassVar[dict[str, str]] = {'token_json': 'GOOGLE_TOKEN_JSON', 'scopes': 'GOOGLE_SCOPES'}
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._credentials: Any = None
+        self._duckdb_token: str | None = None
+        self._services: dict[tuple[str, str], Any] = {}
+
+    def setup(self, secrets: SecretStore) -> None:
+        google_credentials = self.import_module('google.oauth2.credentials')
+        token_document = json.loads(self.require_secret('GOOGLE_TOKEN_JSON'))
+        scopes = self.secret('GOOGLE_SCOPES')
+        self._credentials = google_credentials.Credentials.from_authorized_user_info(
+            token_document,
+            scopes=scopes.split(',') if scopes else list(DEFAULT_SCOPES),
+        )
+
+    # -- credentials -------------------------------------------------------
+    @property
+    def credentials(self) -> Any:
+        """Valid credentials, refreshed on demand."""
+        creds = self._credentials
+        if creds is None:
+            raise RuntimeError('google component is not set up')
+        if not creds.valid and creds.refresh_token:
+            from google.auth.transport.requests import Request
+
+            creds.refresh(Request())
+            self._services.clear()
+        return creds
+
+    def service(self, api: str, version: str) -> Any:
+        """Cached ``googleapiclient`` service, rebuilt after a token refresh."""
+        credentials = self.credentials
+        key = (api, version)
+        if key not in self._services:
+            discovery = self.import_module('googleapiclient.discovery')
+            self._services[key] = discovery.build(api, version, credentials=credentials, cache_discovery=False)
+        return self._services[key]
+
+    # -- duckdb ------------------------------------------------------------
+    def register_duckdb_secret(self) -> None:
+        """(Re)register the DuckDB ``gsheet`` secret with the current access token.
+
+        DuckDB's ``CREATE SECRET`` takes literals only, so the token is inlined;
+        it is quoted defensively and it is an access token that expires in an
+        hour, never the refresh token.
+        """
+        token = self.credentials.token
+        if token == self._duckdb_token:
+            return
+        connection = get_duckdb()
+        connection.execute('INSTALL gsheets FROM community; LOAD gsheets;')
+        connection.execute(
+            f'CREATE OR REPLACE SECRET arrows_gsheet '
+            f'(TYPE gsheet, PROVIDER access_token, TOKEN {quote_literal(token)});'
+        )
+        self._duckdb_token = token
+
+    def health_check(self) -> HealthStatus:
+        try:
+            profile = self.service('oauth2', 'v2').userinfo().get().execute()
+            return HealthStatus(self.name, True, f'user={profile.get("email", "?")}')
+        except Exception as exc:
+            return HealthStatus(self.name, False, f'{type(exc).__name__}: {exc}')
+
+    def close(self) -> None:
+        # Teardown is best effort: the connection may already be gone.
+        with contextlib.suppress(Exception):
+            get_duckdb().execute('DROP SECRET IF EXISTS arrows_gsheet;')
+        self._services.clear()
+        self._credentials = None
+        self._duckdb_token = None
+
+
+COMPONENT = GoogleComponent
+
+
+def token_file_hint(credentials_dir: str | Path = '~/.credentials') -> str:
+    path = Path(credentials_dir).expanduser() / 'google_token.json'
+    return f'Place the OAuth token document at {path} (chmod 600) or set GOOGLE_TOKEN_JSON.'