|
| 1 | +"""Airtable Workspace Component. |
| 2 | +
|
| 3 | +Auto-enumerates Airtable bases × tables via the Meta API and emits one |
| 4 | +Dagster asset per table. Materializing an asset reads records via the |
| 5 | +Records API and emits a DataFrame with the flattened fields. |
| 6 | +""" |
| 7 | +import hashlib |
| 8 | +import json |
| 9 | +from dataclasses import dataclass, field |
| 10 | +from pathlib import Path |
| 11 | +from typing import Any, List, Optional |
| 12 | + |
| 13 | +import dagster as dg |
| 14 | + |
| 15 | +try: |
| 16 | + from dagster.components.component.state_backed_component import StateBackedComponent |
| 17 | + from dagster.components.utils.defs_state import ( |
| 18 | + DefsStateConfig, |
| 19 | + DefsStateConfigArgs, |
| 20 | + ResolvedDefsStateConfig, |
| 21 | + ) |
| 22 | + _HAS_STATE_BACKED = True |
| 23 | +except ImportError: |
| 24 | + StateBackedComponent = None # type: ignore |
| 25 | + DefsStateConfig = None # type: ignore |
| 26 | + DefsStateConfigArgs = None # type: ignore |
| 27 | + ResolvedDefsStateConfig = Any # type: ignore |
| 28 | + _HAS_STATE_BACKED = False |
| 29 | + |
| 30 | + |
| 31 | +@dataclass |
| 32 | +class AirtableTableSelector(dg.Resolvable): |
| 33 | + """Selector for filtering Airtable tables. Fivetran-shape. |
| 34 | +
|
| 35 | + Matches on `<base_name>.<table_name>` OR just `<table_name>` depending on |
| 36 | + whether patterns include a dot. |
| 37 | + """ |
| 38 | + by_name: Optional[List[str]] = None |
| 39 | + by_pattern: Optional[List[str]] = None |
| 40 | + exclude_by_name: Optional[List[str]] = None |
| 41 | + exclude_by_pattern: Optional[List[str]] = None |
| 42 | + |
| 43 | + def matches(self, qualified_name: str) -> bool: |
| 44 | + """qualified_name = 'BaseName.TableName' — selectors can match either.""" |
| 45 | + import fnmatch |
| 46 | + table_only = qualified_name.split(".", 1)[-1] if "." in qualified_name else qualified_name |
| 47 | + candidates = [qualified_name, table_only] |
| 48 | + |
| 49 | + if self.exclude_by_name and any(n in self.exclude_by_name for n in candidates): |
| 50 | + return False |
| 51 | + if self.exclude_by_pattern: |
| 52 | + for p in self.exclude_by_pattern: |
| 53 | + if any(fnmatch.fnmatch(n, p) for n in candidates): |
| 54 | + return False |
| 55 | + if not self.by_name and not self.by_pattern: |
| 56 | + return True |
| 57 | + if self.by_name and any(n in self.by_name for n in candidates): |
| 58 | + return True |
| 59 | + if self.by_pattern: |
| 60 | + for p in self.by_pattern: |
| 61 | + if any(fnmatch.fnmatch(n, p) for n in candidates): |
| 62 | + return True |
| 63 | + return False |
| 64 | + |
| 65 | + |
| 66 | +def _enumerate_airtable(pat, verify_ssl) -> dict: |
| 67 | + import requests |
| 68 | + session = requests.Session() |
| 69 | + session.verify = verify_ssl |
| 70 | + headers = {"Authorization": f"Bearer {pat}", "Accept": "application/json"} |
| 71 | + |
| 72 | + out: dict = {"bases": []} |
| 73 | + try: |
| 74 | + r = session.get("https://api.airtable.com/v0/meta/bases", headers=headers, timeout=60) |
| 75 | + r.raise_for_status() |
| 76 | + body = r.json() or {} |
| 77 | + for b in body.get("bases", []): |
| 78 | + base_id = b.get("id") |
| 79 | + base_name = b.get("name") or base_id |
| 80 | + if not base_id: |
| 81 | + continue |
| 82 | + tables: list = [] |
| 83 | + try: |
| 84 | + tr = session.get( |
| 85 | + f"https://api.airtable.com/v0/meta/bases/{base_id}/tables", |
| 86 | + headers=headers, timeout=30, |
| 87 | + ) |
| 88 | + tr.raise_for_status() |
| 89 | + tbody = tr.json() or {} |
| 90 | + for t in tbody.get("tables", []): |
| 91 | + tid = t.get("id") |
| 92 | + tname = t.get("name") or tid |
| 93 | + if tid and tname: |
| 94 | + tables.append({"id": tid, "name": tname}) |
| 95 | + except Exception: # noqa: BLE001 |
| 96 | + pass |
| 97 | + out["bases"].append({"id": base_id, "name": base_name, "tables": tables}) |
| 98 | + except Exception: # noqa: BLE001 |
| 99 | + pass |
| 100 | + return out |
| 101 | + |
| 102 | + |
| 103 | +if _HAS_STATE_BACKED: |
| 104 | + |
| 105 | + @dataclass |
| 106 | + class AirtableWorkspaceComponent(StateBackedComponent, dg.Resolvable): |
| 107 | + pat_env_var: str # Airtable personal access token |
| 108 | + verify_ssl: bool = True |
| 109 | + |
| 110 | + table_selector: Optional[AirtableTableSelector] = None |
| 111 | + page_size: int = 100 |
| 112 | + |
| 113 | + group_name: Optional[str] = None |
| 114 | + asset_key_prefix: List[str] = field(default_factory=lambda: ["airtable"]) |
| 115 | + compute_kind: str = "airtable" |
| 116 | + |
| 117 | + defs_state: Optional[ResolvedDefsStateConfig] = None |
| 118 | + |
| 119 | + @property |
| 120 | + def defs_state_config(self) -> "DefsStateConfig": |
| 121 | + _key = f"AirtableWorkspace[{hashlib.sha256(self.pat_env_var.encode()).hexdigest()[:12]}]" |
| 122 | + return DefsStateConfig.from_args(self.defs_state, default_key=_key) |
| 123 | + |
| 124 | + def write_state_to_path(self, state_path: Path) -> None: |
| 125 | + import os |
| 126 | + pat = os.environ.get(self.pat_env_var, "") |
| 127 | + snapshot = _enumerate_airtable(pat, self.verify_ssl) |
| 128 | + if self.table_selector is not None: |
| 129 | + for b in snapshot["bases"]: |
| 130 | + b["tables"] = [ |
| 131 | + t for t in b["tables"] |
| 132 | + if self.table_selector.matches(f"{b['name']}.{t['name']}") |
| 133 | + ] |
| 134 | + state_path.write_text(json.dumps(snapshot, indent=2)) |
| 135 | + |
| 136 | + def build_defs_from_state( |
| 137 | + self, context: dg.ComponentLoadContext, state_path: Optional[Path], |
| 138 | + ) -> dg.Definitions: |
| 139 | + if state_path is None or not state_path.exists(): |
| 140 | + return dg.Definitions() |
| 141 | + state = json.loads(state_path.read_text()) |
| 142 | + assets = [] |
| 143 | + for b in state.get("bases", []): |
| 144 | + for t in b.get("tables", []): |
| 145 | + assets.append(self._build_asset(b["id"], b["name"], t["id"], t["name"])) |
| 146 | + return dg.Definitions(assets=assets) |
| 147 | + |
| 148 | + def _build_asset(self, base_id: str, base_name: str, table_id: str, table_name: str): |
| 149 | + _self = self |
| 150 | + safe_base = "".join(c if c.isalnum() or c == "_" else "_" for c in base_name)[:40] or "base" |
| 151 | + safe_table = "".join(c if c.isalnum() or c == "_" else "_" for c in table_name)[:40] or "table" |
| 152 | + key = dg.AssetKey([*self.asset_key_prefix, safe_base, safe_table]) |
| 153 | + |
| 154 | + @dg.asset( |
| 155 | + key=key, group_name=self.group_name, compute_kind=self.compute_kind, |
| 156 | + metadata={ |
| 157 | + "airtable_base_id": dg.MetadataValue.text(base_id), |
| 158 | + "airtable_base_name": dg.MetadataValue.text(base_name), |
| 159 | + "airtable_table_id": dg.MetadataValue.text(table_id), |
| 160 | + "airtable_table_name": dg.MetadataValue.text(table_name), |
| 161 | + }, |
| 162 | + ) |
| 163 | + def _asset(context: dg.AssetExecutionContext): |
| 164 | + import os |
| 165 | + try: |
| 166 | + import pandas as pd |
| 167 | + import requests |
| 168 | + except ImportError as e: |
| 169 | + raise Exception("pandas or requests library not installed") from e |
| 170 | + |
| 171 | + pat = os.environ.get(_self.pat_env_var, "") |
| 172 | + session = requests.Session() |
| 173 | + session.verify = _self.verify_ssl |
| 174 | + headers = {"Authorization": f"Bearer {pat}", "Accept": "application/json"} |
| 175 | + |
| 176 | + url = f"https://api.airtable.com/v0/{base_id}/{table_id}" |
| 177 | + r = session.get(url, headers=headers, params={"pageSize": _self.page_size}, timeout=120) |
| 178 | + r.raise_for_status() |
| 179 | + body = r.json() or {} |
| 180 | + records = body.get("records", []) |
| 181 | + rows = [] |
| 182 | + for rec in records: |
| 183 | + row = {"id": rec.get("id"), "createdTime": rec.get("createdTime")} |
| 184 | + row.update(rec.get("fields") or {}) |
| 185 | + rows.append(row) |
| 186 | + df = pd.DataFrame(rows) |
| 187 | + context.add_output_metadata({ |
| 188 | + "row_count": len(df), |
| 189 | + "airtable_table": f"{base_name}.{table_name}", |
| 190 | + }) |
| 191 | + return df |
| 192 | + |
| 193 | + return _asset |
| 194 | + |
| 195 | +else: |
| 196 | + class AirtableWorkspaceComponent: # type: ignore |
| 197 | + def __init__(self, *args, **kwargs): |
| 198 | + raise ImportError("AirtableWorkspaceComponent requires Dagster with StateBackedComponent support.") |
0 commit comments