local-bigquery 0.3.0__py3-none-any.whl

local_bigquery/__init__.py

@@ -0,0 +1,3 @@
+from local_bigquery.app import app
+
+__all__ = ["app"]

local_bigquery/api/__init__.py

@@ -0,0 +1,30 @@
+import json
+
+from fastapi import APIRouter, Response
+
+from local_bigquery.resource import Resource
+
+
+class Router(APIRouter):
+    def add_api_route(self, path, endpoint, **kwargs):
+        super().add_api_route(
+            path, endpoint, **kwargs | {"response_model_exclude_unset": True}
+        )
+
+
+def paginate(
+    items: list, max_results: int | None, page_token: str | None
+) -> tuple[list, str | None]:
+    start = int(page_token or 0)
+    end = len(items) if max_results is None else start + max_results
+    return items[start:end], str(end) if end < len(items) else None
+
+
+def with_rows(payload: dict, rows: list[str]) -> Response:
+    body = json.dumps(
+        {key: value for key, value in payload.items() if value is not None},
+        default=Resource.dump,
+    )
+    if rows:
+        body = f'{body[:-1]}, "rows": [{",".join(rows)}]}}'
+    return Response(body, media_type="application/json")

local_bigquery/api/datasets.py

@@ -0,0 +1,59 @@
+from fastapi import Header
+
+from local_bigquery.api import Router, paginate
+from local_bigquery.catalog import datasets
+from local_bigquery.models import Dataset, DatasetList, DatasetListDatasetsItem
+
+router = Router(tags=["datasets"])
+SUMMARY_FIELDS = set(DatasetListDatasetsItem.model_fields)
+
+
+@router.get("/projects/{project_id}/datasets")
+def list_datasets(
+    project_id: str,
+    maxResults: int | None = None,
+    pageToken: str | None = None,
+    all: bool = False,
+    filter: str | None = None,
+) -> DatasetList:
+    summaries = [
+        dataset.model_dump(include=SUMMARY_FIELDS, exclude_none=True)
+        for dataset in datasets.list_(project_id, filter, all)
+    ]
+    page, token = paginate(summaries, maxResults, pageToken)
+    return DatasetList(kind="bigquery#datasetList", datasets=page, nextPageToken=token)
+
+
+@router.post("/projects/{project_id}/datasets")
+def insert_dataset(project_id: str, body: Dataset) -> Dataset:
+    return datasets.create(project_id, body)
+
+
+@router.get("/projects/{project_id}/datasets/{dataset_id}")
+def get_dataset(project_id: str, dataset_id: str) -> Dataset:
+    return datasets.get(project_id, dataset_id)
+
+
+@router.patch("/projects/{project_id}/datasets/{dataset_id}")
+def patch_dataset(
+    project_id: str,
+    dataset_id: str,
+    body: Dataset,
+    if_match: str | None = Header(None),
+) -> Dataset:
+    return datasets.update(project_id, dataset_id, body, if_match, replace=False)
+
+
+@router.put("/projects/{project_id}/datasets/{dataset_id}")
+def update_dataset(
+    project_id: str,
+    dataset_id: str,
+    body: Dataset,
+    if_match: str | None = Header(None),
+) -> Dataset:
+    return datasets.update(project_id, dataset_id, body, if_match, replace=True)
+
+
+@router.delete("/projects/{project_id}/datasets/{dataset_id}", status_code=204)
+def delete_dataset(project_id: str, dataset_id: str, deleteContents: bool = False):
+    datasets.delete(project_id, dataset_id, deleteContents)

local_bigquery/api/iam.py

@@ -0,0 +1,29 @@
+from local_bigquery.api import Router
+from local_bigquery.catalog import iam
+from local_bigquery.models import (
+    Policy,
+    SetIamPolicyRequest,
+    TestIamPermissionsRequest,
+    TestIamPermissionsResponse,
+)
+
+router = Router(tags=["iam"])
+RESOURCE = "/{resource:path}"
+
+
+@router.post(f"{RESOURCE}:getIamPolicy")
+def get_policy(resource: str) -> Policy:
+    return iam.get(resource)
+
+
+@router.post(f"{RESOURCE}:setIamPolicy")
+def set_policy(resource: str, body: SetIamPolicyRequest) -> Policy:
+    return iam.set_(resource, body.policy or Policy())
+
+
+@router.post(f"{RESOURCE}:testIamPermissions")
+def test_permissions(
+    resource: str, body: TestIamPermissionsRequest
+) -> TestIamPermissionsResponse:
+    iam.get(resource)
+    return TestIamPermissionsResponse(permissions=body.permissions or [])

local_bigquery/api/jobs.py

@@ -0,0 +1,176 @@
+import uuid
+
+from fastapi import Query
+
+from local_bigquery.api import Router, paginate, with_rows
+from local_bigquery.catalog import tabledata, tables
+from local_bigquery.jobs import runner, store
+from local_bigquery.models import (
+    Job,
+    JobCancelResponse,
+    JobConfiguration,
+    JobList,
+    JobListJobsItem,
+    QueryRequest,
+    TableSchema,
+)
+
+router = Router(tags=["jobs"])
+LIST_FIELDS = set(JobListJobsItem.model_fields)
+QUERY_CONFIGURATION = ("dryRun", "labels", "jobTimeoutMs", "jobCreationMode")
+QUERY_ONLY = ("kind", "formatOptions", "timeoutMs")
+
+
+def results(
+    job: Job, max_results: int | None, start: int, int64_timestamps: bool
+) -> tuple[dict, list[str]]:
+    payload = {"jobReference": job.jobReference, "jobComplete": False}
+    if job.status.state != "DONE":
+        return payload, []
+    if error := runner.error(job):
+        raise error
+    statistics = job.statistics.query
+    payload |= {
+        "jobComplete": True,
+        "cacheHit": False,
+        "totalBytesProcessed": "0",
+        "statementType": statistics.statementType,
+        "numDmlAffectedRows": statistics.numDmlAffectedRows,
+        "dmlStats": statistics.dmlStats,
+        "totalRows": "0",
+    }
+    destination = job.configuration.query.destinationTable
+    if not destination:
+        return payload | {"schema": statistics.schema_}, []
+    page, schema = tabledata.list_rows(
+        *tables.reference(destination), max_results, start, None, int64_timestamps
+    )
+    payload |= {
+        "schema": TableSchema(fields=schema),
+        "totalRows": str(page.total),
+        "pageToken": page.next_token,
+    }
+    return payload, page.rows
+
+
+@router.get("/projects/{project_id}/jobs")
+def list_jobs(
+    project_id: str,
+    maxResults: int | None = None,
+    pageToken: str | None = None,
+    stateFilter: list[str] | None = Query(None),
+    parentJobId: str | None = None,
+    minCreationTime: int | None = None,
+    maxCreationTime: int | None = None,
+) -> JobList:
+    states = [state.upper() for state in stateFilter] if stateFilter else None
+    jobs = [
+        JobListJobsItem.model_validate(
+            job.model_dump(include=LIST_FIELDS, exclude_none=True)
+            | {"state": job.status.state, "errorResult": job.status.errorResult}
+        )
+        for job in store.list_(
+            project_id, states, parentJobId, minCreationTime, maxCreationTime
+        )
+    ]
+    page, token = paginate(jobs, maxResults, pageToken)
+    return JobList(kind="bigquery#jobList", jobs=page, nextPageToken=token)
+
+
+@router.post("/projects/{project_id}/jobs")
+def insert_job(project_id: str, body: Job) -> Job:
+    job_id = (body.jobReference and body.jobReference.jobId) or str(uuid.uuid4())
+    job = runner.submit(project_id, job_id, body.configuration or JobConfiguration())
+    if job.jobReference.jobId is None:
+        return job
+    return runner.submitted(project_id, job_id)
+
+
+@router.get("/projects/{project_id}/jobs/{job_id}")
+def get_job(project_id: str, job_id: str) -> Job:
+    return runner.get(project_id, job_id)
+
+
+@router.post("/projects/{project_id}/jobs/{job_id}/cancel")
+def cancel_job(project_id: str, job_id: str) -> JobCancelResponse:
+    job = runner.cancel(project_id, job_id)
+    return JobCancelResponse(kind="bigquery#jobCancelResponse", job=job)
+
+
+@router.delete("/projects/{project_id}/jobs/{job_id}/delete")
+def delete_job(project_id: str, job_id: str) -> dict:
+    job = runner.get(project_id, job_id)
+    config = job.configuration.query
+    destination = config and config.destinationTable
+    if destination and destination.datasetId == tables.RESULTS:
+        tables.delete(project_id, tables.RESULTS, job_id)
+    store.delete(project_id, job_id)
+    return {}
+
+
+@router.post("/projects/{project_id}/queries")
+def run_query(project_id: str, body: QueryRequest):
+    request = body.given()
+    configuration = JobConfiguration.model_validate(
+        {key: request[key] for key in QUERY_CONFIGURATION if key in request}
+        | {
+            "query": {
+                key: value
+                for key, value in request.items()
+                if key not in (*QUERY_CONFIGURATION, *QUERY_ONLY)
+            }
+        }
+    )
+    job_id = body.requestId or str(uuid.uuid4())
+    job = runner.submit(project_id, job_id, configuration)
+    if job.jobReference.jobId is None:
+        statistics = job.statistics.query
+        return with_rows(
+            {
+                "kind": "bigquery#queryResponse",
+                "jobReference": job.jobReference,
+                "jobComplete": True,
+                "schema": statistics.schema_,
+                "totalBytesProcessed": "0",
+                "statementType": statistics.statementType,
+            },
+            [],
+        )
+    job = runner.wait(project_id, job_id, body.timeoutMs)
+    if error := runner.error(job):
+        raise runner.synchronous(error)
+    int64_timestamps = bool(body.formatOptions and body.formatOptions.useInt64Timestamp)
+    payload, rows = results(job, body.maxResults, 0, int64_timestamps)
+    statistics = job.statistics
+    return with_rows(
+        {"kind": "bigquery#queryResponse"}
+        | payload
+        | {
+            "queryId": job_id,
+            "location": "US",
+            "creationTime": statistics.creationTime,
+            "startTime": statistics.startTime,
+            "endTime": statistics.endTime,
+            "jobCreationReason": job.jobCreationReason,
+            "sessionInfo": statistics.sessionInfo,
+            "totalBytesBilled": "0",
+            "totalSlotMs": "0",
+        },
+        rows,
+    )
+
+
+@router.get("/projects/{project_id}/queries/{job_id}")
+def get_query_results(
+    project_id: str,
+    job_id: str,
+    maxResults: int | None = None,
+    pageToken: str | None = None,
+    startIndex: int = 0,
+    timeoutMs: int | None = None,
+    int64_timestamps: bool = Query(False, alias="formatOptions.useInt64Timestamp"),
+):
+    job = runner.wait(project_id, job_id, timeoutMs)
+    start = int(pageToken) if pageToken else startIndex
+    payload, rows = results(job, maxResults, start, int64_timestamps)
+    return with_rows({"kind": "bigquery#getQueryResultsResponse"} | payload, rows)

local_bigquery/api/models.py

@@ -0,0 +1,40 @@
+from fastapi import Header
+
+from local_bigquery.api import Router, paginate
+from local_bigquery.catalog import models
+from local_bigquery.models import ListModelsResponse, Model
+
+router = Router(tags=["models"])
+MODEL = "/projects/{project_id}/datasets/{dataset_id}/models/{model_id}"
+
+
+@router.get("/projects/{project_id}/datasets/{dataset_id}/models")
+def list_models(
+    project_id: str,
+    dataset_id: str,
+    maxResults: int | None = None,
+    pageToken: str | None = None,
+) -> ListModelsResponse:
+    page, token = paginate(models.list_(project_id, dataset_id), maxResults, pageToken)
+    return ListModelsResponse(models=page, nextPageToken=token)
+
+
+@router.get(MODEL)
+def get_model(project_id: str, dataset_id: str, model_id: str) -> Model:
+    return models.get(project_id, dataset_id, model_id)
+
+
+@router.patch(MODEL)
+def patch_model(
+    project_id: str,
+    dataset_id: str,
+    model_id: str,
+    body: Model,
+    if_match: str | None = Header(None),
+) -> Model:
+    return models.update(project_id, dataset_id, model_id, body, if_match)
+
+
+@router.delete(MODEL, status_code=204)
+def delete_model(project_id: str, dataset_id: str, model_id: str):
+    models.delete(project_id, dataset_id, model_id)

local_bigquery/api/projects.py

@@ -0,0 +1,35 @@
+from local_bigquery.api import Router, paginate
+from local_bigquery.engine import database
+from local_bigquery.models import GetServiceAccountResponse, ProjectList
+
+router = Router(tags=["projects"])
+
+
+@router.get("/projects")
+def list_projects(
+    maxResults: int | None = None, pageToken: str | None = None
+) -> ProjectList:
+    project_ids = database.projects()
+    page, token = paginate(project_ids, maxResults, pageToken)
+    return ProjectList(
+        kind="bigquery#projectList",
+        projects=[
+            {
+                "kind": "bigquery#project",
+                "id": project_id,
+                "numericId": str(abs(hash(project_id))),
+                "friendlyName": project_id,
+                "projectReference": {"projectId": project_id},
+            }
+            for project_id in page
+        ],
+        nextPageToken=token,
+        totalItems=len(project_ids),
+    )
+
+
+@router.get("/projects/{project_id}/serviceAccount")
+def get_service_account(project_id: str) -> GetServiceAccountResponse:
+    return GetServiceAccountResponse(
+        email=f"service-account@{project_id}.iam.gserviceaccount.com"
+    )

local_bigquery/api/routines.py

@@ -0,0 +1,81 @@
+from fastapi import Header
+
+from local_bigquery.api import Router, paginate
+from local_bigquery.catalog import datasets, metadata, names, routines
+from local_bigquery.errors import BigQueryError, already_exists
+from local_bigquery.jobs.query import run_ddl
+from local_bigquery.models import ListRoutinesResponse, Routine
+
+router = Router(tags=["routines"])
+ROUTINES = "/projects/{project_id}/datasets/{dataset_id}/routines"
+ROUTINE = f"{ROUTINES}/{{routine_id}}"
+
+
+def load(project_id: str, dataset_id: str, routine_id: str) -> Routine:
+    datasets.get(project_id, dataset_id)
+    return routines.get(project_id, dataset_id, routine_id)
+
+
+def define(
+    project_id: str, dataset_id: str, routine_id: str, body: Routine, replace: bool
+) -> Routine:
+    reference = {
+        "projectId": project_id,
+        "datasetId": dataset_id,
+        "routineId": routine_id,
+    }
+    body = body.replace(routineReference=reference)
+    run_ddl(project_id, routines.statement(body, replace))
+    stored = routines.load(project_id, dataset_id, routine_id)
+    return routines.save(
+        project_id, dataset_id, routine_id, stored.replace(**body.dump())
+    )
+
+
+@router.get(ROUTINES)
+def list_routines(
+    project_id: str,
+    dataset_id: str,
+    maxResults: int | None = None,
+    pageToken: str | None = None,
+) -> ListRoutinesResponse:
+    datasets.get(project_id, dataset_id)
+    page, token = paginate(
+        routines.list_(project_id, dataset_id), maxResults, pageToken
+    )
+    return ListRoutinesResponse(routines=page, nextPageToken=token)
+
+
+@router.post(ROUTINES)
+def insert_routine(project_id: str, dataset_id: str, body: Routine) -> Routine:
+    datasets.get(project_id, dataset_id)
+    routine_id = body.routineReference and body.routineReference.routineId
+    if routines.load(project_id, dataset_id, routine_id) is not None:
+        raise already_exists("Routine", names.label(project_id, dataset_id, routine_id))
+    return define(project_id, dataset_id, routine_id, body, False)
+
+
+@router.get(ROUTINE)
+def get_routine(project_id: str, dataset_id: str, routine_id: str) -> Routine:
+    return load(project_id, dataset_id, routine_id)
+
+
+@router.put(ROUTINE)
+def update_routine(
+    project_id: str,
+    dataset_id: str,
+    routine_id: str,
+    body: Routine,
+    if_match: str | None = Header(None),
+) -> Routine:
+    if not body.routineType:
+        raise BigQueryError("invalid", "Routine type must be specified")
+    current = load(project_id, dataset_id, routine_id)
+    metadata.check_etag(current, if_match)
+    return define(project_id, dataset_id, routine_id, body, True)
+
+
+@router.delete(ROUTINE, status_code=204)
+def delete_routine(project_id: str, dataset_id: str, routine_id: str):
+    kind = routines.ROUTINE_TYPES[load(project_id, dataset_id, routine_id).routineType]
+    run_ddl(project_id, f"DROP {kind} `{project_id}.{dataset_id}.{routine_id}`")

local_bigquery/api/row_access_policies.py

@@ -0,0 +1,73 @@
+from local_bigquery.api import Router, paginate
+from local_bigquery.catalog import row_access
+from local_bigquery.models import (
+    BatchDeleteRowAccessPoliciesRequest,
+    ListRowAccessPoliciesResponse,
+    RowAccessPolicy,
+)
+
+router = Router(tags=["rowAccessPolicies"])
+POLICIES = (
+    "/projects/{project_id}/datasets/{dataset_id}/tables/{table_id}/rowAccessPolicies"
+)
+
+
+@router.get(POLICIES)
+def list_policies(
+    project_id: str,
+    dataset_id: str,
+    table_id: str,
+    pageSize: int | None = None,
+    pageToken: str | None = None,
+) -> ListRowAccessPoliciesResponse:
+    policies = row_access.list_(project_id, dataset_id, table_id)
+    page, token = paginate(policies, pageSize, pageToken)
+    return ListRowAccessPoliciesResponse(rowAccessPolicies=page, nextPageToken=token)
+
+
+@router.post(POLICIES)
+def insert_policy(
+    project_id: str, dataset_id: str, table_id: str, body: RowAccessPolicy
+) -> RowAccessPolicy:
+    return row_access.save(project_id, dataset_id, table_id, body)
+
+
+@router.post(f"{POLICIES}:batchDelete")
+def batch_delete(
+    project_id: str,
+    dataset_id: str,
+    table_id: str,
+    body: BatchDeleteRowAccessPoliciesRequest,
+) -> dict:
+    row_access.delete(project_id, dataset_id, table_id, *body.policyIds or [])
+    return {}
+
+
+@router.get(f"{POLICIES}/{{policy_id}}")
+def get_policy(
+    project_id: str, dataset_id: str, table_id: str, policy_id: str
+) -> RowAccessPolicy:
+    return row_access.get(project_id, dataset_id, table_id, policy_id)
+
+
+@router.put(f"{POLICIES}/{{policy_id}}")
+def update_policy(
+    project_id: str,
+    dataset_id: str,
+    table_id: str,
+    policy_id: str,
+    body: RowAccessPolicy,
+) -> RowAccessPolicy:
+    row_access.get(project_id, dataset_id, table_id, policy_id)
+    body = body.merged(
+        RowAccessPolicy(rowAccessPolicyReference={"policyId": policy_id})
+    )
+    return row_access.save(project_id, dataset_id, table_id, body, replace=True)
+
+
+@router.delete(f"{POLICIES}/{{policy_id}}")
+def delete_policy(
+    project_id: str, dataset_id: str, table_id: str, policy_id: str
+) -> dict:
+    row_access.delete(project_id, dataset_id, table_id, policy_id)
+    return {}

local_bigquery/api/tables.py

@@ -0,0 +1,115 @@
+from fastapi import Header, Query
+
+from local_bigquery.api import Router, paginate, with_rows
+from local_bigquery.catalog import tabledata, tables
+from local_bigquery.jobs.query import translate_view
+from local_bigquery.models import (
+    Table,
+    TableDataInsertAllRequest,
+    TableDataInsertAllResponse,
+    TableList,
+)
+
+router = Router(tags=["tables"])
+TABLE = "/projects/{project_id}/datasets/{dataset_id}/tables/{table_id}"
+
+
+@router.get("/projects/{project_id}/datasets/{dataset_id}/tables")
+def list_tables(
+    project_id: str,
+    dataset_id: str,
+    maxResults: int | None = None,
+    pageToken: str | None = None,
+) -> TableList:
+    summaries = tables.list_(project_id, dataset_id)
+    page, token = paginate(summaries, maxResults, pageToken)
+    return TableList(
+        kind="bigquery#tableList",
+        tables=[table.dump() for table in page],
+        nextPageToken=token,
+        totalItems=len(summaries),
+    )
+
+
+@router.post("/projects/{project_id}/datasets/{dataset_id}/tables")
+def insert_table(project_id: str, dataset_id: str, body: Table) -> Table:
+    return tables.create(project_id, dataset_id, body, translate_view)
+
+
+@router.get(TABLE)
+def get_table(
+    project_id: str,
+    dataset_id: str,
+    table_id: str,
+    selectedFields: str | None = None,
+    view: str | None = None,
+) -> Table:
+    return tables.get(project_id, dataset_id, table_id, selectedFields, view)
+
+
+@router.patch(TABLE)
+def patch_table(
+    project_id: str,
+    dataset_id: str,
+    table_id: str,
+    body: Table,
+    if_match: str | None = Header(None),
+) -> Table:
+    return tables.update(project_id, dataset_id, table_id, body, if_match, False)
+
+
+@router.put(TABLE)
+def update_table(
+    project_id: str,
+    dataset_id: str,
+    table_id: str,
+    body: Table,
+    if_match: str | None = Header(None),
+) -> Table:
+    return tables.update(project_id, dataset_id, table_id, body, if_match, True)
+
+
+@router.delete(TABLE, status_code=204)
+def delete_table(project_id: str, dataset_id: str, table_id: str):
+    tables.delete(project_id, dataset_id, table_id)
+
+
+@router.post(f"{TABLE}/insertAll")
+def insert_all(
+    project_id: str, dataset_id: str, table_id: str, body: TableDataInsertAllRequest
+) -> TableDataInsertAllResponse:
+    response = tabledata.insert_all(project_id, dataset_id, table_id, body)
+    return TableDataInsertAllResponse(
+        kind="bigquery#tableDataInsertAllResponse", **response
+    )
+
+
+@router.get(f"{TABLE}/data")
+def list_rows(
+    project_id: str,
+    dataset_id: str,
+    table_id: str,
+    maxResults: int | None = None,
+    pageToken: str | None = None,
+    startIndex: int = 0,
+    selectedFields: str | None = None,
+    int64_timestamps: bool = Query(False, alias="formatOptions.useInt64Timestamp"),
+):
+    start = int(pageToken) if pageToken else startIndex
+    page, _ = tabledata.list_rows(
+        project_id,
+        dataset_id,
+        table_id,
+        maxResults,
+        start,
+        selectedFields,
+        int64_timestamps,
+    )
+    return with_rows(
+        {
+            "kind": "bigquery#tableDataList",
+            "totalRows": str(page.total),
+            "pageToken": page.next_token,
+        },
+        page.rows,
+    )