local-data-studio 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (89) hide show
  1. local_data_studio/__init__.py +3 -0
  2. local_data_studio/__main__.py +6 -0
  3. local_data_studio/app.py +150 -0
  4. local_data_studio/cli.py +154 -0
  5. local_data_studio/runtime_config.py +77 -0
  6. local_data_studio/server/__init__.py +1 -0
  7. local_data_studio/server/api/__init__.py +1 -0
  8. local_data_studio/server/api/analysis.py +123 -0
  9. local_data_studio/server/api/atlas_proxy.py +156 -0
  10. local_data_studio/server/api/datasets.py +176 -0
  11. local_data_studio/server/api/jobs.py +231 -0
  12. local_data_studio/server/api/mutations.py +63 -0
  13. local_data_studio/server/api/schemas.py +117 -0
  14. local_data_studio/server/api/services.py +99 -0
  15. local_data_studio/server/api/static.py +34 -0
  16. local_data_studio/server/atlas.py +91 -0
  17. local_data_studio/server/atlas_cache.py +20 -0
  18. local_data_studio/server/atlas_cache_patch.py +40 -0
  19. local_data_studio/server/atlas_components/__init__.py +1 -0
  20. local_data_studio/server/atlas_components/command.py +80 -0
  21. local_data_studio/server/atlas_components/contracts.py +101 -0
  22. local_data_studio/server/atlas_components/dataset.py +233 -0
  23. local_data_studio/server/atlas_components/embedding_backends.py +260 -0
  24. local_data_studio/server/atlas_components/embedding_session.py +89 -0
  25. local_data_studio/server/atlas_components/image_values.py +193 -0
  26. local_data_studio/server/atlas_components/images.py +63 -0
  27. local_data_studio/server/atlas_components/output_frames.py +157 -0
  28. local_data_studio/server/atlas_components/ports.py +140 -0
  29. local_data_studio/server/atlas_components/process.py +175 -0
  30. local_data_studio/server/atlas_components/projection.py +223 -0
  31. local_data_studio/server/atlas_components/projection_inputs.py +242 -0
  32. local_data_studio/server/atlas_components/prompts.py +93 -0
  33. local_data_studio/server/atlas_components/readiness.py +60 -0
  34. local_data_studio/server/atlas_components/reducers.py +116 -0
  35. local_data_studio/server/atlas_components/runtime.py +205 -0
  36. local_data_studio/server/atlas_components/sampling.py +73 -0
  37. local_data_studio/server/atlas_components/service.py +221 -0
  38. local_data_studio/server/cache.py +143 -0
  39. local_data_studio/server/column_stats/__init__.py +35 -0
  40. local_data_studio/server/column_stats/accumulator.py +162 -0
  41. local_data_studio/server/column_stats/heuristics.py +187 -0
  42. local_data_studio/server/column_stats/service.py +52 -0
  43. local_data_studio/server/config.py +269 -0
  44. local_data_studio/server/dataset_readers/__init__.py +5 -0
  45. local_data_studio/server/dataset_readers/common.py +216 -0
  46. local_data_studio/server/dataset_readers/contracts.py +38 -0
  47. local_data_studio/server/dataset_readers/delimited.py +174 -0
  48. local_data_studio/server/dataset_readers/json_reader.py +86 -0
  49. local_data_studio/server/dataset_readers/jsonl.py +182 -0
  50. local_data_studio/server/dataset_readers/line.py +39 -0
  51. local_data_studio/server/dataset_readers/line_cursor.py +75 -0
  52. local_data_studio/server/dataset_readers/line_indexing.py +57 -0
  53. local_data_studio/server/dataset_readers/parquet.py +202 -0
  54. local_data_studio/server/dataset_readers/service.py +228 -0
  55. local_data_studio/server/db.py +187 -0
  56. local_data_studio/server/delete_ops.py +261 -0
  57. local_data_studio/server/deleted_rows.py +30 -0
  58. local_data_studio/server/eda.py +106 -0
  59. local_data_studio/server/eda_reports.py +178 -0
  60. local_data_studio/server/embedder_capabilities.py +384 -0
  61. local_data_studio/server/embedder_models.py +70 -0
  62. local_data_studio/server/files.py +190 -0
  63. local_data_studio/server/jobs.py +284 -0
  64. local_data_studio/server/line_index.py +126 -0
  65. local_data_studio/server/llm.py +40 -0
  66. local_data_studio/server/llm_client.py +91 -0
  67. local_data_studio/server/llm_profiles.py +297 -0
  68. local_data_studio/server/llm_prompt.py +81 -0
  69. local_data_studio/server/llm_service.py +39 -0
  70. local_data_studio/server/readers.py +29 -0
  71. local_data_studio/server/serialization.py +118 -0
  72. local_data_studio/server/sql.py +277 -0
  73. local_data_studio/server/stats.py +33 -0
  74. local_data_studio/static/app/application.js +1789 -0
  75. local_data_studio/static/app/atlas.js +354 -0
  76. local_data_studio/static/app/dom.js +77 -0
  77. local_data_studio/static/app/formatting.js +69 -0
  78. local_data_studio/static/app/http.js +74 -0
  79. local_data_studio/static/app/images.js +212 -0
  80. local_data_studio/static/app/llm.js +129 -0
  81. local_data_studio/static/app/state.js +58 -0
  82. local_data_studio/static/app.js +1 -0
  83. local_data_studio/static/index.html +384 -0
  84. local_data_studio/static/styles.css +1568 -0
  85. local_data_studio-1.0.0.dist-info/METADATA +683 -0
  86. local_data_studio-1.0.0.dist-info/RECORD +89 -0
  87. local_data_studio-1.0.0.dist-info/WHEEL +4 -0
  88. local_data_studio-1.0.0.dist-info/entry_points.txt +2 -0
  89. local_data_studio-1.0.0.dist-info/licenses/LICENSE +21 -0
@@ -0,0 +1,3 @@
1
+ """Local Data Studio package."""
2
+
3
+ __version__ = "0.1.0"
@@ -0,0 +1,6 @@
1
+ """Run Local Data Studio with ``python -m local_data_studio``."""
2
+
3
+ from .cli import main
4
+
5
+ if __name__ == "__main__":
6
+ main()
@@ -0,0 +1,150 @@
1
+ """FastAPI application assembly for Local Data Studio."""
2
+
3
+ import asyncio
4
+ from collections.abc import AsyncIterator
5
+ from contextlib import asynccontextmanager
6
+
7
+ from fastapi import FastAPI
8
+
9
+ from .server.api.analysis import column_stats, count_rows, llm_models, nl_query, run_eda, run_query, search
10
+ from .server.api.analysis import router as analysis_router
11
+ from .server.api.atlas_proxy import router as atlas_proxy_router
12
+ from .server.api.atlas_proxy import stop_atlas_instance
13
+ from .server.api.datasets import (
14
+ column_sample,
15
+ embedder_models,
16
+ get_config,
17
+ get_schema,
18
+ list_files,
19
+ preview,
20
+ raw_file,
21
+ raw_row,
22
+ upload_files,
23
+ )
24
+ from .server.api.datasets import (
25
+ router as datasets_router,
26
+ )
27
+ from .server.api.jobs import (
28
+ cancel_job,
29
+ get_job,
30
+ start_atlas_job,
31
+ start_atlas_query_job,
32
+ start_count_job,
33
+ start_eda_job,
34
+ start_eda_query_job,
35
+ start_index_job,
36
+ start_query_job,
37
+ start_search_job,
38
+ start_stats_job,
39
+ )
40
+ from .server.api.jobs import (
41
+ router as jobs_router,
42
+ )
43
+ from .server.api.mutations import delete_column, delete_row
44
+ from .server.api.mutations import router as mutations_router
45
+ from .server.api.schemas import (
46
+ AtlasQueryRequest,
47
+ AtlasRequest,
48
+ CountJobRequest,
49
+ DeleteColumnRequest,
50
+ DeleteRowRequest,
51
+ EdaQueryRequest,
52
+ EdaRequest,
53
+ IndexJobRequest,
54
+ NLQueryRequest,
55
+ QueryRequest,
56
+ RawRowRequest,
57
+ SearchJobRequest,
58
+ StatsJobRequest,
59
+ )
60
+ from .server.api.static import NoCacheStaticFiles, mount_static_files
61
+ from .server.atlas_components.runtime import AtlasRuntime
62
+ from .server.config import ATLAS_MAX_INSTANCES
63
+ from .server.jobs import JobStore
64
+
65
+
66
+ def create_app(*, job_store: JobStore | None = None, atlas_runtime: AtlasRuntime | None = None) -> FastAPI:
67
+ """Assemble the application with API routes before static catch-all mounts.
68
+
69
+ Returns:
70
+ A new application instance. Callers own the instance and may add routes.
71
+ """
72
+ owned_job_store = job_store or JobStore()
73
+ owned_atlas_runtime = atlas_runtime or AtlasRuntime(max_instances=ATLAS_MAX_INSTANCES)
74
+
75
+ @asynccontextmanager
76
+ async def lifespan(_: FastAPI) -> AsyncIterator[None]:
77
+ await owned_atlas_runtime.start_proxy_client()
78
+ try:
79
+ yield
80
+ finally:
81
+ owned_atlas_runtime.begin_shutdown()
82
+ owned_job_store.begin_shutdown()
83
+ await asyncio.to_thread(owned_job_store.wait_for_idle, 10.0)
84
+ await asyncio.to_thread(owned_atlas_runtime.terminate_all)
85
+ await owned_atlas_runtime.close_proxy_client()
86
+ owned_job_store.shutdown(wait_timeout=0.0)
87
+
88
+ application = FastAPI(title="Data Viewer", lifespan=lifespan)
89
+ application.state.job_store = owned_job_store
90
+ application.state.atlas_runtime = owned_atlas_runtime
91
+ application.include_router(datasets_router)
92
+ application.include_router(analysis_router)
93
+ application.include_router(jobs_router)
94
+ application.include_router(mutations_router)
95
+ application.include_router(atlas_proxy_router)
96
+ mount_static_files(application)
97
+ return application
98
+
99
+
100
+ app = create_app()
101
+
102
+
103
+ __all__ = [
104
+ "AtlasQueryRequest",
105
+ "AtlasRequest",
106
+ "CountJobRequest",
107
+ "DeleteColumnRequest",
108
+ "DeleteRowRequest",
109
+ "EdaQueryRequest",
110
+ "EdaRequest",
111
+ "IndexJobRequest",
112
+ "NLQueryRequest",
113
+ "NoCacheStaticFiles",
114
+ "QueryRequest",
115
+ "RawRowRequest",
116
+ "SearchJobRequest",
117
+ "StatsJobRequest",
118
+ "app",
119
+ "cancel_job",
120
+ "column_sample",
121
+ "column_stats",
122
+ "count_rows",
123
+ "create_app",
124
+ "delete_column",
125
+ "delete_row",
126
+ "embedder_models",
127
+ "get_config",
128
+ "get_job",
129
+ "get_schema",
130
+ "list_files",
131
+ "llm_models",
132
+ "nl_query",
133
+ "preview",
134
+ "raw_file",
135
+ "raw_row",
136
+ "run_eda",
137
+ "run_query",
138
+ "search",
139
+ "start_atlas_job",
140
+ "start_atlas_query_job",
141
+ "start_count_job",
142
+ "start_eda_job",
143
+ "start_eda_query_job",
144
+ "start_index_job",
145
+ "start_query_job",
146
+ "start_search_job",
147
+ "start_stats_job",
148
+ "stop_atlas_instance",
149
+ "upload_files",
150
+ ]
@@ -0,0 +1,154 @@
1
+ """Command-line entrypoint for Local Data Studio."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ import os
7
+ from collections.abc import Mapping, Sequence
8
+ from pathlib import Path
9
+ from typing import Any
10
+
11
+ import uvicorn
12
+
13
+ from .runtime_config import CONFIG_FILE_ENV, apply_settings_environment, config_section, read_runtime_config
14
+
15
+ PATH_ENV_KEYS: dict[str, str] = {
16
+ "workspace_dir": "LOCAL_DATA_STUDIO_WORKSPACE_DIR",
17
+ "env_file": "LOCAL_DATA_STUDIO_ENV_FILE",
18
+ "data_dir": "DATA_DIR",
19
+ "data_file": "DATA_FILE",
20
+ "cache_dir": "CACHE_DIR",
21
+ "models_dir": "EMBEDDER_MODELS_DIR",
22
+ "file_serve_roots": "FILE_SERVE_ROOTS",
23
+ "vis_exclude_dirs": "VIS_EXCLUDE_DIRS",
24
+ "vis_exclude_files": "VIS_EXCLUDE_FILES",
25
+ }
26
+
27
+
28
+ def _set_path_env(name: str, value: str | None, *, overwrite: bool = True) -> None:
29
+ if value:
30
+ if overwrite or name not in os.environ:
31
+ os.environ[name] = str(Path(value).expanduser())
32
+
33
+
34
+ def _resolve_path(value: str, base_dir: Path) -> str:
35
+ path = Path(value).expanduser()
36
+ if not path.is_absolute():
37
+ path = base_dir / path
38
+ return str(path.resolve())
39
+
40
+
41
+ def _coerce_path_setting(value: Any, base_dir: Path) -> str | None:
42
+ if value is None:
43
+ return None
44
+ if isinstance(value, str):
45
+ stripped = value.strip()
46
+ return _resolve_path(stripped, base_dir) if stripped else None
47
+ if isinstance(value, Sequence) and not isinstance(value, (bytes, bytearray, str)):
48
+ resolved = [_resolve_path(str(item), base_dir) for item in value if str(item).strip()]
49
+ return ",".join(resolved) if resolved else None
50
+ return str(value)
51
+
52
+
53
+ def _workspace_from(args: argparse.Namespace, config: Mapping[str, Any], config_path: Path | None) -> Path:
54
+ paths = config_section(config, "paths")
55
+ raw_workspace = args.workspace_dir or os.environ.get("LOCAL_DATA_STUDIO_WORKSPACE_DIR") or paths.get("workspace_dir")
56
+ if raw_workspace:
57
+ base_dir = config_path.parent if config_path else Path.cwd()
58
+ return Path(_resolve_path(str(raw_workspace), base_dir))
59
+ return Path.cwd().resolve()
60
+
61
+
62
+ def _apply_config_environment(config: Mapping[str, Any], workspace_dir: Path) -> None:
63
+ paths = config_section(config, "paths")
64
+ _set_path_env("LOCAL_DATA_STUDIO_WORKSPACE_DIR", str(workspace_dir), overwrite=False)
65
+ for key, env_name in PATH_ENV_KEYS.items():
66
+ if key == "workspace_dir":
67
+ continue
68
+ coerced = _coerce_path_setting(paths.get(key), workspace_dir)
69
+ _set_path_env(env_name, coerced, overwrite=False)
70
+
71
+
72
+ def configure_runtime_environment(args: argparse.Namespace) -> tuple[str, int, bool]:
73
+ """Resolve CLI, environment, and TOML settings before importing the app.
74
+
75
+ Explicit CLI paths take precedence over environment and config values. Relative
76
+ config paths are resolved from the config file; other defaults use the workspace.
77
+
78
+ Returns:
79
+ The host, port, and reload values passed to Uvicorn.
80
+
81
+ Raises:
82
+ OSError: The requested config file cannot be read.
83
+ tomllib.TOMLDecodeError: The config file is not valid TOML.
84
+ """
85
+ config, config_path = read_runtime_config(args.config)
86
+ if config_path:
87
+ os.environ[CONFIG_FILE_ENV] = str(config_path)
88
+ workspace_dir = _workspace_from(args, config, config_path)
89
+ _apply_config_environment(config, workspace_dir)
90
+ apply_settings_environment(config)
91
+
92
+ _set_path_env("LOCAL_DATA_STUDIO_WORKSPACE_DIR", str(workspace_dir), overwrite=bool(args.workspace_dir))
93
+ _set_path_env("LOCAL_DATA_STUDIO_ENV_FILE", args.env_file)
94
+ _set_path_env("DATA_DIR", args.data_dir)
95
+ _set_path_env("DATA_FILE", args.data_file)
96
+ _set_path_env("CACHE_DIR", args.cache_dir)
97
+ _set_path_env("EMBEDDER_MODELS_DIR", args.models_dir)
98
+ if args.file_serve_roots:
99
+ os.environ["FILE_SERVE_ROOTS"] = ",".join(str(Path(item).expanduser()) for item in args.file_serve_roots)
100
+
101
+ server = config_section(config, "server")
102
+ host = args.host or os.environ.get("LOCAL_DATA_STUDIO_HOST") or str(server.get("host") or "127.0.0.1")
103
+ port = args.port or int(os.environ.get("LOCAL_DATA_STUDIO_PORT") or server.get("port") or 8000)
104
+ reload = bool(args.reload if args.reload is not None else _coerce_bool(os.environ.get("LOCAL_DATA_STUDIO_RELOAD"), server.get("reload", False)))
105
+ return host, port, reload
106
+
107
+
108
+ def _coerce_bool(value: Any, default: bool = False) -> bool:
109
+ if value is None:
110
+ return default
111
+ if isinstance(value, bool):
112
+ return value
113
+ normalized = str(value).strip().lower()
114
+ if normalized in {"1", "true", "yes", "on"}:
115
+ return True
116
+ if normalized in {"0", "false", "no", "off"}:
117
+ return False
118
+ return default
119
+
120
+
121
+ def build_parser() -> argparse.ArgumentParser:
122
+ """Build the command-line parser without reading runtime configuration."""
123
+ parser = argparse.ArgumentParser(description="Run the Local Data Studio web app.")
124
+ parser.add_argument("--config", help="Path to a Local Data Studio TOML config file.")
125
+ parser.add_argument("--workspace-dir", help="Base directory for .env, data, cache, and models/embedder defaults.")
126
+ parser.add_argument("--env-file", help="Path to the .env file to load.")
127
+ parser.add_argument("--host", help="Host interface for the local web server.")
128
+ parser.add_argument("--port", type=int, help="Port for the local web server.")
129
+ parser.add_argument("--cache-dir", help="Directory for Local Data Studio cache files.")
130
+ parser.add_argument("--models-dir", help="Directory containing HuggingFace encoder model directories.")
131
+ parser.add_argument("--file-serve-roots", nargs="+", help="Directories from which local image files may be served.")
132
+ parser.add_argument("--reload", action=argparse.BooleanOptionalAction, default=None, help="Enable uvicorn reload for local development.")
133
+
134
+ data_source = parser.add_mutually_exclusive_group()
135
+ data_source.add_argument("--data-dir", help="Directory to search for datasets.")
136
+ data_source.add_argument("--data-file", help="Single dataset file to open.")
137
+ return parser
138
+
139
+
140
+ def main(argv: Sequence[str] | None = None) -> None:
141
+ """Configure the workspace and run the Uvicorn web server.
142
+
143
+ Args:
144
+ argv: Arguments without the executable name. ``None`` reads ``sys.argv``.
145
+ """
146
+ args = build_parser().parse_args(argv)
147
+ host, port, reload = configure_runtime_environment(args)
148
+
149
+ uvicorn.run(
150
+ "local_data_studio.app:app",
151
+ host=host,
152
+ port=port,
153
+ reload=reload,
154
+ )
@@ -0,0 +1,77 @@
1
+ """Shared loading helpers for Local Data Studio TOML configuration."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ import tomllib
7
+ from collections.abc import Mapping
8
+ from pathlib import Path
9
+ from typing import Any
10
+
11
+ CONFIG_FILE_ENV = "LOCAL_DATA_STUDIO_CONFIG_FILE"
12
+
13
+ # These TOML keys mirror environment-backed Settings fields. Paths remain in
14
+ # [paths], while [settings] owns the operational EDA, Atlas, and deletion
15
+ # controls so a checked local TOML can replace a large .env file.
16
+ SETTINGS_ENV_KEYS: dict[str, str] = {
17
+ "eda_row_limit": "EDA_ROW_LIMIT",
18
+ "allow_delete_data": "ALLOW_DELETE_DATA",
19
+ "eda_cell_max_chars": "EDA_CELL_MAX_CHARS",
20
+ "eda_nested_policy": "EDA_NESTED_POLICY",
21
+ "eda_cache_max_bytes": "EDA_CACHE_MAX_BYTES",
22
+ "atlas_host": "ATLAS_HOST",
23
+ "atlas_port": "ATLAS_PORT",
24
+ "atlas_max_instances": "ATLAS_MAX_INSTANCES",
25
+ "atlas_sample": "ATLAS_SAMPLE",
26
+ "atlas_batch_size": "ATLAS_BATCH_SIZE",
27
+ "atlas_cache_max_bytes": "ATLAS_CACHE_MAX_BYTES",
28
+ "atlas_text_max_chars": "ATLAS_TEXT_MAX_CHARS",
29
+ "atlas_embedding_dtype": "ATLAS_EMBEDDING_DTYPE",
30
+ "atlas_umap_projection_mode": "ATLAS_UMAP_PROJECTION_MODE",
31
+ "atlas_umap_anchor_sample": "ATLAS_UMAP_ANCHOR_SAMPLE",
32
+ "atlas_trust_remote_code": "ATLAS_TRUST_REMOTE_CODE",
33
+ }
34
+
35
+
36
+ def read_runtime_config(path: str | None) -> tuple[dict[str, Any], Path | None]:
37
+ """Read an optional TOML configuration file.
38
+
39
+ Args:
40
+ path: Explicit path, or ``None`` to use ``LOCAL_DATA_STUDIO_CONFIG_FILE``.
41
+
42
+ Returns:
43
+ The parsed root table and resolved source path. Both are empty when no
44
+ configuration file is selected.
45
+
46
+ Raises:
47
+ OSError: The selected file cannot be read.
48
+ tomllib.TOMLDecodeError: The selected file is not valid TOML.
49
+ """
50
+ selected = path or os.environ.get(CONFIG_FILE_ENV)
51
+ if not selected:
52
+ return {}, None
53
+ config_path = Path(selected).expanduser().resolve()
54
+ with config_path.open("rb") as handle:
55
+ loaded = tomllib.load(handle)
56
+ return loaded, config_path
57
+
58
+
59
+ def config_section(config: Mapping[str, Any], name: str) -> Mapping[str, Any]:
60
+ """Return a mapping-valued root section, or an empty mapping."""
61
+ value = config.get(name)
62
+ return value if isinstance(value, Mapping) else {}
63
+
64
+
65
+ def apply_settings_environment(config: Mapping[str, Any]) -> None:
66
+ """Expose `[settings]` values through the existing Settings environment API.
67
+
68
+ Existing OS environment variables are never replaced. This preserves the
69
+ documented precedence over TOML while allowing Pydantic Settings to keep
70
+ its existing validation and `.env` fallback behavior.
71
+ """
72
+ settings = config_section(config, "settings")
73
+ for key, environment_name in SETTINGS_ENV_KEYS.items():
74
+ value = settings.get(key)
75
+ if value is None or environment_name in os.environ:
76
+ continue
77
+ os.environ[environment_name] = str(value).lower() if isinstance(value, bool) else str(value)
@@ -0,0 +1 @@
1
+ """Backend package."""
@@ -0,0 +1 @@
1
+ """HTTP API routers for Local Data Studio."""
@@ -0,0 +1,123 @@
1
+ """Synchronous analysis and SQL endpoints."""
2
+
3
+ from typing import Any
4
+
5
+ from fastapi import APIRouter, HTTPException, Query
6
+
7
+ from ..config import DEFAULT_LIMIT, DEFAULT_SAMPLE
8
+ from ..db import (
9
+ build_table_response,
10
+ count_relation_rows,
11
+ describe_relation,
12
+ fetch_rows_with_rowid,
13
+ normalize_pagination,
14
+ open_connection,
15
+ quote_ident,
16
+ relation_sql,
17
+ relation_with_rowid_sql,
18
+ )
19
+ from ..deleted_rows import deleted_row_ids_for
20
+ from ..files import resolve_data_file
21
+ from ..llm import generate_sql_request, public_llm_models
22
+ from ..sql import execute_query_guarded
23
+ from ..stats import compute_column_stats
24
+ from .schemas import EdaRequest, NLQueryRequest, QueryRequest
25
+ from .services import eda_reports_service, reject_large_sync_operation
26
+
27
+ router = APIRouter()
28
+
29
+
30
+ @router.get("/api/llm_models")
31
+ def llm_models() -> dict[str, Any]:
32
+ """Return server-managed model profiles without secrets or endpoints."""
33
+ return public_llm_models()
34
+
35
+
36
+ @router.get("/api/column_stats")
37
+ def column_stats(file: str = Query(...), sample: int | None = Query(DEFAULT_SAMPLE)) -> dict[str, Any]:
38
+ """Return bounded sample statistics, rejecting large synchronous scans."""
39
+ path = resolve_data_file(file)
40
+ reject_large_sync_operation(path, "synchronous column stats")
41
+ return compute_column_stats(file, path, sample if sample is not None else DEFAULT_SAMPLE)
42
+
43
+
44
+ @router.post("/api/eda")
45
+ def run_eda(payload: EdaRequest) -> dict[str, Any]:
46
+ """Generate or reuse an EDA report and return its cache URL."""
47
+ path = resolve_data_file(payload.file)
48
+ return eda_reports_service().generate_dataset_eda_report(
49
+ file_name=payload.file,
50
+ path=path,
51
+ sample=payload.sample,
52
+ mode=payload.mode,
53
+ force=payload.force,
54
+ )
55
+
56
+
57
+ @router.post("/api/nl_query")
58
+ def nl_query(payload: NLQueryRequest) -> dict[str, Any]:
59
+ """Generate SQL from a non-empty natural-language prompt and dataset schema."""
60
+ path = resolve_data_file(payload.file)
61
+ prompt = payload.prompt.strip()
62
+ if not prompt:
63
+ raise HTTPException(status_code=400, detail="prompt must not be empty")
64
+ relation, params = relation_sql(path)
65
+ with open_connection() as connection:
66
+ columns = describe_relation(connection, relation, params)
67
+ result = generate_sql_request(prompt, columns, payload.sample, payload.model)
68
+ return {"sql": result.sql, "model": result.model_id, "model_label": result.model_label}
69
+
70
+
71
+ @router.get("/api/count")
72
+ def count_rows(file: str = Query(...)) -> dict[str, Any]:
73
+ """Count rows synchronously only when the dataset is below the safety limit."""
74
+ path = resolve_data_file(file)
75
+ reject_large_sync_operation(path, "synchronous row count")
76
+ return {"file": file, "count": count_relation_rows(path)}
77
+
78
+
79
+ @router.get("/api/search")
80
+ def search(
81
+ file: str = Query(...),
82
+ query: str = Query(...),
83
+ limit: int | None = Query(DEFAULT_LIMIT),
84
+ offset: int | None = Query(0),
85
+ ) -> dict[str, Any]:
86
+ """Search textual columns synchronously with bounded pagination.
87
+
88
+ Raises:
89
+ HTTPException: The query is blank or the dataset requires a background job.
90
+ """
91
+ path = resolve_data_file(file)
92
+ reject_large_sync_operation(path, "synchronous search")
93
+ deleted_ids = deleted_row_ids_for(path)
94
+ relation, params = relation_with_rowid_sql(path, deleted_ids)
95
+ search_term = query.strip()
96
+ if not search_term:
97
+ raise HTTPException(status_code=400, detail="query must not be empty")
98
+ limit_value, offset_value = normalize_pagination(limit, offset)
99
+
100
+ with open_connection() as connection:
101
+ base_relation, base_params = relation_sql(path)
102
+ columns = describe_relation(connection, base_relation, base_params)
103
+ text_columns = [column["name"] for column in columns if "CHAR" in column["type"].upper() or "TEXT" in column["type"].upper()]
104
+ if not text_columns:
105
+ return build_table_response(file, [column["name"] for column in columns], [], limit_value, offset_value, [])
106
+ like_clauses = " OR ".join(f"CAST({quote_ident(column)} AS VARCHAR) ILIKE ?" for column in text_columns)
107
+ values = params + [f"%{search_term}%"] * len(text_columns)
108
+ search_sql = f"SELECT * FROM ({relation}) WHERE {like_clauses} LIMIT {limit_value} OFFSET {offset_value}"
109
+ result_columns, rows, row_ids = fetch_rows_with_rowid(connection, search_sql, values)
110
+ return build_table_response(file, result_columns, rows, limit_value, offset_value, row_ids)
111
+
112
+
113
+ @router.post("/api/query")
114
+ def run_query(payload: QueryRequest) -> dict[str, Any]:
115
+ """Run validated read-only SQL under configured resource limits."""
116
+ path = resolve_data_file(payload.file)
117
+ return execute_query_guarded(
118
+ file_name=payload.file,
119
+ path=path,
120
+ sql=payload.sql,
121
+ limit=payload.limit,
122
+ offset=payload.offset,
123
+ )
@@ -0,0 +1,156 @@
1
+ """Same-origin streaming reverse proxy for registered Atlas instances."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import re
6
+ from collections.abc import AsyncIterator, Iterable
7
+ from urllib.parse import urljoin, urlsplit, urlunsplit
8
+
9
+ import httpx
10
+ from fastapi import APIRouter, HTTPException, Request
11
+ from starlette.background import BackgroundTask
12
+ from starlette.responses import RedirectResponse, StreamingResponse
13
+
14
+ from ..atlas_components.runtime import AtlasInstance, AtlasRuntime, AtlasRuntimeError
15
+
16
+ router = APIRouter()
17
+
18
+ INSTANCE_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]{32,64}$")
19
+ HOP_BY_HOP_HEADERS = {
20
+ b"connection",
21
+ b"keep-alive",
22
+ b"proxy-authenticate",
23
+ b"proxy-authorization",
24
+ b"te",
25
+ b"trailer",
26
+ b"transfer-encoding",
27
+ b"upgrade",
28
+ }
29
+ PRIVATE_REQUEST_HEADERS = {b"authorization", b"cookie", b"host", b"content-length"}
30
+ PRIVATE_RESPONSE_HEADERS = {b"set-cookie"}
31
+ PROXY_METHODS = ["GET", "HEAD", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"]
32
+
33
+
34
+ def _runtime(request: Request) -> AtlasRuntime:
35
+ return request.app.state.atlas_runtime
36
+
37
+
38
+ def _connection_tokens(headers: Iterable[tuple[bytes, bytes]]) -> set[bytes]:
39
+ tokens: set[bytes] = set()
40
+ for name, value in headers:
41
+ if name.lower() == b"connection":
42
+ tokens.update(token.strip().lower() for token in value.split(b",") if token.strip())
43
+ return tokens
44
+
45
+
46
+ def _filtered_headers(
47
+ headers: Iterable[tuple[bytes, bytes]],
48
+ *,
49
+ private: set[bytes],
50
+ ) -> list[tuple[bytes, bytes]]:
51
+ raw_headers = list(headers)
52
+ blocked = HOP_BY_HOP_HEADERS | private | _connection_tokens(raw_headers)
53
+ return [(name, value) for name, value in raw_headers if name.lower() not in blocked]
54
+
55
+
56
+ def _resolve_instance(runtime: AtlasRuntime, instance_id: str) -> AtlasInstance:
57
+ if INSTANCE_ID_PATTERN.fullmatch(instance_id) is None:
58
+ raise HTTPException(status_code=404, detail="Atlas instance not found")
59
+ instance = runtime.resolve(instance_id)
60
+ if instance is None:
61
+ raise HTTPException(status_code=404, detail="Atlas instance not found")
62
+ return instance
63
+
64
+
65
+ def _upstream_url(request: Request, instance: AtlasInstance) -> httpx.URL:
66
+ raw_path = request.scope.get("raw_path")
67
+ if not isinstance(raw_path, bytes):
68
+ raw_path = request.url.path.encode("utf-8")
69
+ prefix = f"/atlas/{instance.instance_id}".encode()
70
+ if not raw_path.startswith(prefix):
71
+ raise HTTPException(status_code=404, detail="Atlas instance not found")
72
+ target_path = raw_path[len(prefix) :] or b"/"
73
+ if not target_path.startswith(b"/"):
74
+ raise HTTPException(status_code=400, detail="Invalid Atlas proxy path")
75
+ query = request.scope.get("query_string", b"")
76
+ raw_target = target_path + (b"?" + query if query else b"")
77
+ # The host is intentionally constant: request data and child output cannot
78
+ # turn this route into a general-purpose server-side request proxy.
79
+ return httpx.URL(scheme="http", host="127.0.0.1", port=instance.port, raw_path=raw_target)
80
+
81
+
82
+ def _rewrite_location(value: bytes, *, upstream_url: httpx.URL, instance: AtlasInstance) -> bytes:
83
+ location = value.decode("latin-1")
84
+ resolved = urljoin(str(upstream_url), location)
85
+ parsed = urlsplit(resolved)
86
+ upstream = urlsplit(str(upstream_url))
87
+ if (parsed.scheme, parsed.hostname, parsed.port) != (upstream.scheme, upstream.hostname, upstream.port):
88
+ return value
89
+ path = parsed.path.lstrip("/")
90
+ proxy_path = instance.proxy_path + path
91
+ rewritten = urlunsplit(("", "", proxy_path, parsed.query, parsed.fragment))
92
+ return rewritten.encode("latin-1")
93
+
94
+
95
+ async def _response_body(upstream: httpx.Response) -> AsyncIterator[bytes]:
96
+ try:
97
+ async for chunk in upstream.aiter_raw():
98
+ yield chunk
99
+ finally:
100
+ await upstream.aclose()
101
+
102
+
103
+ @router.delete("/api/atlas/instances/{instance_id}")
104
+ def stop_atlas_instance(instance_id: str, request: Request) -> dict[str, object]:
105
+ """Stop one live Atlas child selected by its opaque instance identifier."""
106
+ runtime = _runtime(request)
107
+ _resolve_instance(runtime, instance_id)
108
+ if not runtime.stop(instance_id):
109
+ raise HTTPException(status_code=404, detail="Atlas instance not found")
110
+ return {"instance_id": instance_id, "stopped": True}
111
+
112
+
113
+ @router.api_route("/atlas/{instance_id}", methods=PROXY_METHODS, include_in_schema=False)
114
+ async def redirect_atlas_root(instance_id: str, request: Request) -> RedirectResponse:
115
+ """Canonicalize Atlas roots so relative frontend assets resolve correctly."""
116
+ _resolve_instance(_runtime(request), instance_id)
117
+ query = request.scope.get("query_string", b"")
118
+ suffix = f"?{query.decode('latin-1')}" if query else ""
119
+ return RedirectResponse(f"/atlas/{instance_id}/{suffix}", status_code=307)
120
+
121
+
122
+ @router.api_route("/atlas/{instance_id}/{path:path}", methods=PROXY_METHODS, include_in_schema=False)
123
+ async def proxy_atlas(instance_id: str, path: str, request: Request) -> StreamingResponse:
124
+ """Stream one HTTP request to a registered loopback Atlas process."""
125
+ del path # Routing uses the decoded value; forwarding uses ASGI raw_path.
126
+ runtime = _runtime(request)
127
+ instance = _resolve_instance(runtime, instance_id)
128
+ upstream_url = _upstream_url(request, instance)
129
+ request_headers = _filtered_headers(request.headers.raw, private=PRIVATE_REQUEST_HEADERS)
130
+ content = request.stream() if request.method not in {"GET", "HEAD"} else None
131
+ upstream_request = runtime.proxy_client.build_request(
132
+ request.method,
133
+ upstream_url,
134
+ headers=request_headers,
135
+ content=content,
136
+ )
137
+ try:
138
+ upstream = await runtime.proxy_client.send(upstream_request, stream=True)
139
+ except httpx.TimeoutException as exc:
140
+ raise HTTPException(status_code=504, detail="Atlas did not respond before the proxy timeout") from exc
141
+ except (httpx.RequestError, AtlasRuntimeError) as exc:
142
+ raise HTTPException(status_code=502, detail="Atlas is not reachable through the local proxy") from exc
143
+
144
+ response_headers = _filtered_headers(upstream.headers.raw, private=PRIVATE_RESPONSE_HEADERS)
145
+ response_headers = [
146
+ (name, _rewrite_location(value, upstream_url=upstream_url, instance=instance) if name.lower() == b"location" else value)
147
+ for name, value in response_headers
148
+ ]
149
+ response_headers.append((b"referrer-policy", b"no-referrer"))
150
+ response = StreamingResponse(
151
+ _response_body(upstream),
152
+ status_code=upstream.status_code,
153
+ background=BackgroundTask(upstream.aclose),
154
+ )
155
+ response.raw_headers = response_headers
156
+ return response