diff --git a/cirro/helpers/workspace.py b/cirro/helpers/workspace.py new file mode 100644 index 0000000..c075600 --- /dev/null +++ b/cirro/helpers/workspace.py @@ -0,0 +1,86 @@ +""" +Query the project's Sheets with DuckDB from inside a Cirro workspace. + +Only workspaces can query Sheets directly: the read permission belongs to the workspace's AWS role. +Everywhere else, use `cirro.sdk.project.DataPortalProject.query_sheets` or `cirro.sdk.sheet.DataPortalSheet`, +which go through the Cirro API and also cover virtual views, column display names and writes. + +No Cirro login is needed. Refer to a sheet as `sheets..`. + +Requires the `sheets` extra: `pip install "cirro[sheets]"`. + +```python +from cirro.helpers.workspace import connect_sheets, query_sheets + +counts = query_sheets("SELECT icd_code, COUNT(*) FROM sheets.default.diagnoses GROUP BY icd_code") + +conn = connect_sheets() +conn.sql("SELECT * FROM sheets.default.diagnoses").df() +``` +""" +import os +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from duckdb import DuckDBPyConnection + from pandas import DataFrame + + +def connect_sheets() -> 'DuckDBPyConnection': + """ + Open an in-memory DuckDB connection with the project's Sheets attached as the `sheets` catalog. + + Only works inside a Cirro workspace with Sheets enabled. Credentials come from the workspace's AWS + role and are read once, when the connection is made; if a long-running session later fails with an + expired token, call `connect_sheets` again. + + Returns: + `duckdb.DuckDBPyConnection` + """ + # Cirro sets this, with the project, account and region, on workspaces whose role can read Sheets + if os.environ.get("CIRRO_SHEETS_ENABLED") != "true": + raise RuntimeError("connect_sheets() only works inside a Cirro workspace with Sheets enabled. " + "Elsewhere, use project.query_sheets().") + try: + import duckdb + except ImportError: + raise ImportError("The duckdb library is required to query Sheets directly. " + "Please install it using 'pip install \"cirro[sheets]\"'.") + + conn = duckdb.connect() + conn.execute(_attach_sheets_sql(project_id=os.environ["CIRRO_PROJECT_ID"], + account_id=os.environ["CIRRO_AWS_ACCOUNT_ID"], + region=os.environ["AWS_DEFAULT_REGION"])) + return conn + + +def query_sheets(sql: str) -> 'DataFrame': + """ + Run a query against the project's Sheets and return the result as a Pandas DataFrame. + + Only works inside a Cirro workspace with Sheets enabled. Opens a new connection for each call, so it + never holds expired credentials. To run many queries, or to keep tables between them, use + `connect_sheets` instead. + + Args: + sql (str): SQL to run; refer to a sheet as `sheets..
` + + Returns: + `pandas.DataFrame` + """ + with connect_sheets() as conn: + return conn.sql(sql).df() + + +def _attach_sheets_sql(project_id: str, account_id: str, region: str) -> str: + # The same statements the workspace image writes to ~/.cirro/duckdb/sheets.sql for the DuckDB CLI + return f""" + INSTALL httpfs; INSTALL aws; INSTALL iceberg; + LOAD httpfs; LOAD aws; LOAD iceberg; + SET iceberg_unsafe_skip_puffin_verification = true; + CREATE SECRET IF NOT EXISTS aws (TYPE s3, PROVIDER credential_chain); + ATTACH 'arn:aws:s3tables:{region}:{account_id}:bucket/project-{project_id}' AS sheets ( + TYPE iceberg, + ENDPOINT_TYPE s3_tables + ); + """ diff --git a/poetry.lock b/poetry.lock index cb86e40..7f527ee 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1028,6 +1028,55 @@ docs = ["myst-parser (==0.18.0)", "sphinx (==5.1.1)"] ssh = ["paramiko (>=2.4.3)"] websockets = ["websocket-client (>=1.3.0)"] +[[package]] +name = "duckdb" +version = "1.5.6" +description = "DuckDB in-process database" +optional = true +python-versions = ">=3.10.0" +groups = ["main"] +markers = "extra == \"sheets\"" +files = [ + {file = "duckdb-1.5.6-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:64db8a6700e81fe419fba130d8f1780686ad40fbf2eb69f78d2a1533728a0549"}, + {file = "duckdb-1.5.6-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d6d1eac4de11779bb249b89b0544916ad65751da031df5c5f6d779c85b753109"}, + {file = "duckdb-1.5.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:56355a543a79c7f4d8576d27edcbd9aaed19a562a0901188b021c10f4c818800"}, + {file = "duckdb-1.5.6-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:95a6b91bb9149950baeb5d02466c006550d0ea98b9d10f15f7d614a8eb32e174"}, + {file = "duckdb-1.5.6-cp310-cp310-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dbd348e9ebdc8b28f1f9930efb5a74a382063c35d9c43901075566fbae50ab5c"}, + {file = "duckdb-1.5.6-cp310-cp310-win_amd64.whl", hash = "sha256:f14551eef9180fc72869e2d9a2896410a8826169e22495e98a825abaa0eac1a7"}, + {file = "duckdb-1.5.6-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:c88700d0ee68ad149a0cc624df21b0f21efc136ea2449aaadd7cd0c9a564962a"}, + {file = "duckdb-1.5.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:03e4f1b10a8b8ff476eb2b73955590fadbcef978da1167c593114c5edf763960"}, + {file = "duckdb-1.5.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:34623eaabd2c66ba5c20f1a39486321c3b7d32e4e0e001ced95f81e3372dd361"}, + {file = "duckdb-1.5.6-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:56c0f71c6bee982e9c30568bb12371bf66b26bf129c75d8d7f60bc69d6590a2c"}, + {file = "duckdb-1.5.6-cp311-cp311-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:73b108c04c932b36c2fa4e41110cc1c3c8cd510eb49f065f92d050be8e6929fd"}, + {file = "duckdb-1.5.6-cp311-cp311-win_amd64.whl", hash = "sha256:dda311932cf5aae955a53fe28a4fc1700c2ab5fa02dc1f165abdd5ec6c39141e"}, + {file = "duckdb-1.5.6-cp311-cp311-win_arm64.whl", hash = "sha256:df5ae02af278e084f54a9730a9f4f211ed736d0bd8f3bc12af925c2effb5b33d"}, + {file = "duckdb-1.5.6-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:48d07d0651aaeac2c3974afd37599970154b7b79b54c18f27c319c14ccf98d9d"}, + {file = "duckdb-1.5.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:79de3dfa8705b1ba0d59e7e3252e40ff399e0afd12f485502a6c7bf7c2fd809a"}, + {file = "duckdb-1.5.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:dcccce20965e6986cd083fdf192c461685ad0b93cd1ccd0b2a8207f1185f078b"}, + {file = "duckdb-1.5.6-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce89a1025a5317ebe9c520876c48032b5247ac574865486648b1a004f6009875"}, + {file = "duckdb-1.5.6-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bc9619ed7d4ffa117b5155d84b44794366bb6635178d78ed5e13a6024845c757"}, + {file = "duckdb-1.5.6-cp312-cp312-win_amd64.whl", hash = "sha256:09ff51b230219f0d8b47fc8a1e17fb595ba9fab0c3d96a6de4d00b8ff86b3cf1"}, + {file = "duckdb-1.5.6-cp312-cp312-win_arm64.whl", hash = "sha256:b8d795c8b2d5634b3269f974aa97f1fdf878f62f032317a52252a151b693fb1e"}, + {file = "duckdb-1.5.6-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ae352646374cacf48e9981cf031191c494865192fc436d13667a2531fc5d1da3"}, + {file = "duckdb-1.5.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:5a1261e90785e9d29953293e44f60fa073bd1137098924e8de21a037a861b051"}, + {file = "duckdb-1.5.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:97dd7a555b8f5298b76bc7d48a11cb2c64336e8de9bfde783cffb86ea9f54807"}, + {file = "duckdb-1.5.6-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:364992ba1089a2b327391cfcb68fd0bd0ce9090cf293baef861a0ba6847abfee"}, + {file = "duckdb-1.5.6-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:644f54ce99b3b61844bc9a3fe80e0aecb1ea4084b1fffc4396d1569db6111679"}, + {file = "duckdb-1.5.6-cp313-cp313-win_amd64.whl", hash = "sha256:ced693d33ddcee2e5345f077d342c87d2aaa80e41c514e64c9ff2d4e5963c251"}, + {file = "duckdb-1.5.6-cp313-cp313-win_arm64.whl", hash = "sha256:41ecc75bb9328d72d154a705c1a653d2c5c60f686a5c0c6578aa80020753c884"}, + {file = "duckdb-1.5.6-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:aa21d2ad803b2524326e8622d7d96b2bb1ff1d5b60368e1978ee805df9c21fb3"}, + {file = "duckdb-1.5.6-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:8a1b2ad27d414068cbca06c55cfa802eece10f86ea4812ff082f8ab4cb25fc85"}, + {file = "duckdb-1.5.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c79c6d222b1d015cde73b5139087186b00db65357fb4e2c94c2308fbbf465a72"}, + {file = "duckdb-1.5.6-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1052b8050ef5696e2c0d8c836949c72f3dd11f0690466acbea739613e8e2750b"}, + {file = "duckdb-1.5.6-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:19c5e485e59613b8878d1670bcaa7a010f53c5a4da5ae8e08863e5e529ca6182"}, + {file = "duckdb-1.5.6-cp314-cp314-win_amd64.whl", hash = "sha256:ebcbd09cd8578ab1093393e9b16289cda0e8f1791ac595bf00eb5bad75c3cf00"}, + {file = "duckdb-1.5.6-cp314-cp314-win_arm64.whl", hash = "sha256:820a8384faef11cd86068ea48c5da57ce2d8f1c7b3d2bdb9be3398317a7c3728"}, + {file = "duckdb-1.5.6.tar.gz", hash = "sha256:166a91dbfacfc0c9f08cc76c0243cb6d3d4296bfab5bad72a3cfb63140a5b7c8"}, +] + +[package.extras] +all = ["adbc-driver-manager", "fsspec", "ipython", "numpy", "pandas", "pyarrow"] + [[package]] name = "eido" version = "0.2.6" @@ -4547,9 +4596,10 @@ type = ["pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] [extras] nextflow = ["nf-core"] pyodide = ["requests-aws4auth"] +sheets = ["duckdb"] wdl = ["miniwdl"] [metadata] lock-version = "2.1" python-versions = ">3.11.0,<4.0" -content-hash = "b228b256ec6ce562140c3b6da1bc1b4fbe9650c3042394898b7583bb97893357" +content-hash = "3811a13427a5af94fab8b5f3f99fbbd106a22e7585d43b348b72ea22bf4c10c9" diff --git a/pyproject.toml b/pyproject.toml index c5d2f83..92a5e74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ requests-aws4auth = {version="1.3.2", optional=true} miniwdl = {version="^1.13.0", optional=true} referencing = "^0.36.2" nf-core = {version = "^3.3.2", optional=true} +duckdb = {version = ">=1.5.6", optional=true} awscrt = "==0.30.0" [tool.poetry.group.dev.dependencies] @@ -38,6 +39,7 @@ pdoc = "15.0.0" pyodide = ["requests-aws4auth"] wdl = ["miniwdl"] nextflow = ["nf-core"] +sheets = ["duckdb"] [tool.poetry.scripts] cirro-cli = "cirro.cli.cli:main" diff --git a/tests/test_workspace.py b/tests/test_workspace.py new file mode 100644 index 0000000..814f73d --- /dev/null +++ b/tests/test_workspace.py @@ -0,0 +1,44 @@ +import os +import unittest +from unittest.mock import patch + +import duckdb + +from cirro.helpers.workspace import connect_sheets, query_sheets + +_WORKSPACE_ENV = { + "CIRRO_SHEETS_ENABLED": "true", + "CIRRO_PROJECT_ID": "project-1", + "CIRRO_AWS_ACCOUNT_ID": "123456789012", + "AWS_DEFAULT_REGION": "us-west-2" +} + + +class TestConnectSheets(unittest.TestCase): + def _attach_sql(self, env): + with patch.dict(os.environ, env, clear=True), patch("duckdb.connect") as duckdb_connect: + connect_sheets() + return duckdb_connect.return_value.execute.call_args.args[0] + + def test_attaches_project_bucket(self): + sql = self._attach_sql(_WORKSPACE_ENV) + self.assertIn("ATTACH 'arn:aws:s3tables:us-west-2:123456789012:bucket/project-project-1' AS sheets", sql) + + def test_refuses_outside_workspace(self): + with self.assertRaisesRegex(RuntimeError, "only works inside a Cirro workspace"): + self._attach_sql({}) + + def test_refuses_when_sheets_disabled(self): + with self.assertRaisesRegex(RuntimeError, "only works inside a Cirro workspace"): + self._attach_sql({**_WORKSPACE_ENV, "CIRRO_SHEETS_ENABLED": "false"}) + + +class TestQuerySheets(unittest.TestCase): + def test_returns_dataframe_and_closes_connection(self): + conn = duckdb.connect() + with patch("cirro.helpers.workspace.connect_sheets", return_value=conn): + df = query_sheets("SELECT 42 AS answer") + + self.assertEqual(df["answer"].tolist(), [42]) + with self.assertRaises(duckdb.ConnectionException): + conn.sql("SELECT 1")