datajunction-query 0.0.1a65__tar.gz → 0.0.1a66__tar.gz

{datajunction_query-0.0.1a65 → datajunction_query-0.0.1a66}/.coveragerc

@@ -2,7 +2,9 @@
 [run]
 branch = True
 source = dj
-# omit = bad_file.py
+omit =
+    djqs/config.py
+    djqs/exceptions.py
 
 [paths]
 source =

{datajunction_query-0.0.1a65 → datajunction_query-0.0.1a66}/Dockerfile

@@ -2,6 +2,10 @@ FROM jupyter/pyspark-notebook
 USER root
 WORKDIR /code
 COPY . /code
+RUN apt-get update && apt-get install -y \
+    libpq-dev \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
 RUN pip install -e .[uvicorn]
 CMD ["uvicorn", "djqs.api.main:app", "--host", "0.0.0.0", "--port", "8001", "--reload"]
 EXPOSE 8001

{datajunction_query-0.0.1a65 → datajunction_query-0.0.1a66}/Makefile

@@ -27,10 +27,10 @@ docker-run-with-cockroachdb:
 	docker compose -f docker-compose.yml -f docker-compose.cockroachdb.yml up
 
 test: pyenv
-	pdm run pytest --cov=djqs -vv tests/ --doctest-modules djqs --without-integration --without-slow-integration ${PYTEST_ARGS}
+	pdm run pytest --cov=djqs -vv tests/ --cov-report term-missing --doctest-modules djqs --without-integration --without-slow-integration ${PYTEST_ARGS}
 
 integration: pyenv
-	pdm run pytest --cov=djqs -vv tests/ --doctest-modules djqs --with-integration --with-slow-integration
+	pdm run pytest --cov=djqs -vv tests/ --cov-report term-missing --doctest-modules djqs --with-integration --with-slow-integration
 
 clean:
 	pyenv virtualenv-delete djqs

{datajunction_query-0.0.1a65 → datajunction_query-0.0.1a66}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: datajunction-query
-Version: 0.0.1a65
+Version: 0.0.1a66
 Summary: OSS Implementation of a DataJunction Query Service
 Project-URL: repository, https://github.com/DataJunction/dj
 Author-email: DataJunction Authors <roberto@dealmeida.net>
@@ -14,7 +14,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Python: <4.0,>=3.8
+Requires-Python: <4.0,>=3.10
 Requires-Dist: accept-types==0.4.1
 Requires-Dist: cachelib>=0.4.0
 Requires-Dist: duckdb-engine
@@ -22,15 +22,16 @@ Requires-Dist: duckdb==0.8.1
 Requires-Dist: fastapi>=0.79.0
 Requires-Dist: importlib-metadata
 Requires-Dist: msgpack>=1.0.3
+Requires-Dist: psycopg[async,pool]>=3.2.1
+Requires-Dist: pytest-asyncio>=0.24.0
+Requires-Dist: pytest-integration>=0.2.3
 Requires-Dist: python-dotenv==0.19.2
 Requires-Dist: pyyaml>=6.0.1
 Requires-Dist: requests<=2.29.0,>=2.28.2
 Requires-Dist: rich>=10.16.2
 Requires-Dist: snowflake-connector-python>=3.3.1
-Requires-Dist: sqlalchemy-utils>=0.37.7
-Requires-Dist: sqlalchemy<2.0.0,>=1.4.41
-Requires-Dist: sqlmodel<1.0.0,>=0.0.8
-Requires-Dist: sqlparse<1.0.0,>=0.4.3
+Requires-Dist: sqlalchemy>=2.0.34
+Requires-Dist: tenacity>=9.0.0
 Requires-Dist: toml>=0.10.2
 Requires-Dist: trino>=0.324.0
 Provides-Extra: uvicorn

{datajunction_query-0.0.1a65 → datajunction_query-0.0.1a66}/alembic/env.py

@@ -2,18 +2,13 @@
 Environment for Alembic migrations.
 """
 # pylint: disable=no-member, unused-import, no-name-in-module, import-error
-
 from logging.config import fileConfig
 
-from sqlmodel import SQLModel, create_engine
+from sqlalchemy import create_engine
 
 from alembic import context
-from djqs.models.catalog import Catalog
-from djqs.models.engine import Engine
-from djqs.models.query import Query
-from djqs.utils import get_settings
 
-settings = get_settings()
+DEFAULT_URI = "postgresql+psycopg://dj:dj@postgres_metadata:5432/dj"
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
@@ -28,7 +23,7 @@ if config.config_file_name is not None:
 # for 'autogenerate' support
 # from myapp import mymodel
 # target_metadata = mymodel.Base.metadata
-target_metadata = SQLModel.metadata
+target_metadata = None  # pylint: disable=invalid-name
 
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
@@ -48,12 +43,13 @@ def run_migrations_offline():
     script output.
 
     """
-    url = settings.index
+    x_args = context.get_x_argument(as_dictionary=True)
     context.configure(
-        url=url,
+        url=x_args.get("uri") or DEFAULT_URI,
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
+        version_table="alembic_version_djqs",
     )
 
     with context.begin_transaction():
@@ -67,10 +63,15 @@ def run_migrations_online():
     and associate a connection with the context.
 
     """
-    connectable = create_engine(settings.index)
+    x_args = context.get_x_argument(as_dictionary=True)
+    connectable = create_engine(x_args.get("uri") or DEFAULT_URI)
 
     with connectable.connect() as connection:
-        context.configure(connection=connection, target_metadata=target_metadata)
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata,
+            version_table="alembic_version_djqs",
+        )
 
         with context.begin_transaction():
             context.run_migrations()

datajunction_query-0.0.1a66/alembic/versions/2024_09_09_0000-b8f22b3549c7_initial_migration.py

@@ -0,0 +1,41 @@
+"""Initial migration
+
+Revision ID: b8f22b3549c7
+Revises:
+Create Date: 2024-09-09 06:00:00.000000+00:00
+
+"""
+# pylint: disable=no-member, invalid-name, missing-function-docstring, unused-import, no-name-in-module
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b8f22b3549c7"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.execute(
+        """
+    CREATE TABLE query (
+        id UUID PRIMARY KEY,
+        catalog_name VARCHAR NOT NULL,
+        engine_name VARCHAR NOT NULL,
+        engine_version VARCHAR NOT NULL,
+        submitted_query VARCHAR NOT NULL,
+        async_ BOOLEAN NOT NULL,
+        executed_query VARCHAR,
+        scheduled TIMESTAMP,
+        started TIMESTAMP,
+        finished TIMESTAMP,
+        state VARCHAR NOT NULL,
+        progress FLOAT NOT NULL
+    )
+    """,
+    )
+
+
+def downgrade():
+    op.execute("DROP TABLE query")

{datajunction_query-0.0.1a65 → datajunction_query-0.0.1a66}/djqs/__about__.py

@@ -1,4 +1,4 @@
 """
 Version for Hatch
 """
-__version__ = "0.0.1a65"
+__version__ = "0.0.1a66"

{datajunction_query-0.0.1a65 → datajunction_query-0.0.1a66}/djqs/api/helpers.py

@@ -1,48 +1,12 @@
 """
 Helper functions for API
 """
-from http import HTTPStatus
 from typing import Any, Dict, List, Optional
 
-from fastapi import HTTPException
-from sqlalchemy import inspect
-from sqlalchemy.exc import NoResultFound, NoSuchTableError, OperationalError
-from sqlmodel import Session, create_engine, select
+from sqlalchemy import create_engine, inspect
+from sqlalchemy.exc import NoSuchTableError, OperationalError
 
 from djqs.exceptions import DJException, DJTableNotFound
-from djqs.models.catalog import Catalog
-from djqs.models.engine import Engine
-
-
-def get_catalog(session: Session, name: str) -> Catalog:
-    """
-    Get a catalog by name
-    """
-    statement = select(Catalog).where(Catalog.name == name)
-    catalog = session.exec(statement).one_or_none()
-    if not catalog:
-        raise DJException(
-            message=f"Catalog with name `{name}` does not exist.",
-            http_status_code=404,
-        )
-    return catalog
-
-
-def get_engine(session: Session, name: str, version: str) -> Engine:
-    """
-    Return an Engine instance given an engine name and version
-    """
-    statement = (
-        select(Engine).where(Engine.name == name).where(Engine.version == version)
-    )
-    try:
-        engine = session.exec(statement).one()
-    except NoResultFound as exc:
-        raise HTTPException(
-            status_code=HTTPStatus.NOT_FOUND,
-            detail=f"Engine not found: `{name}` version `{version}`",
-        ) from exc
-    return engine
 
 
 def get_columns(

{datajunction_query-0.0.1a65 → datajunction_query-0.0.1a66}/djqs/api/main.py

@@ -12,29 +12,41 @@ from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
+from psycopg.rows import dict_row
+from psycopg_pool import AsyncConnectionPool
 
 from djqs import __version__
-from djqs.api import catalogs, engines, queries, tables
-from djqs.config import load_djqs_config
+from djqs.api import queries, tables
 from djqs.exceptions import DJException
-from djqs.utils import get_session, get_settings
+from djqs.utils import get_settings
 
 _logger = logging.getLogger(__name__)
 
 settings = get_settings()
-session = next(get_session())
 
 
 @asynccontextmanager
-async def lifespan(app: FastAPI):  # pylint: disable=W0621,W0613
+async def lifespan(fastapi_app: FastAPI):
     """
-    Load DJQS config on app startup
+    Create a postgres connection pool and store it in the app state
     """
+    _logger.info("Starting PostgreSQL connection pool...")
+    pool = AsyncConnectionPool(
+        settings.index,
+        kwargs={"row_factory": dict_row},
+        check=AsyncConnectionPool.check_connection,
+        min_size=5,
+        max_size=20,
+        timeout=15,
+    )
+    fastapi_app.state.pool = pool
     try:
-        load_djqs_config(settings=settings, session=session)
-    except Exception as e:  # pylint: disable=W0718,C0103
-        _logger.warning("Could not load DJQS config: %s", e)
-    yield
+        _logger.info("PostgreSQL connection pool started with DSN: %s", settings.index)
+        yield
+    finally:
+        _logger.info("Closing PostgreSQL connection pool")
+        await pool.close()
+        _logger.info("PostgreSQL connection pool closed")
 
 
 app = FastAPI(
@@ -47,12 +59,10 @@ app = FastAPI(
     },
     lifespan=lifespan,
 )
-app.include_router(catalogs.get_router)
-app.include_router(engines.get_router)
 app.include_router(queries.router)
 app.include_router(tables.router)
-app.include_router(catalogs.post_router) if settings.enable_dynamic_config else None
-app.include_router(engines.post_router) if settings.enable_dynamic_config else None
+
+app.router.lifespan_context = lifespan
 
 
 @app.exception_handler(DJException)

datajunction_query-0.0.1a66/djqs/api/queries.py

@@ -0,0 +1,259 @@
+"""
+Query related APIs.
+"""
+
+import json
+import logging
+import uuid
+from dataclasses import asdict
+from http import HTTPStatus
+from typing import Any, Dict, List, Optional
+
+import msgpack
+from accept_types import get_best_match
+from fastapi import (
+    APIRouter,
+    BackgroundTasks,
+    Body,
+    Depends,
+    Header,
+    HTTPException,
+    Request,
+    Response,
+)
+from psycopg_pool import AsyncConnectionPool
+
+from djqs.config import Settings
+from djqs.db.postgres import DBQuery, get_postgres_pool
+from djqs.engine import process_query
+from djqs.models.query import (
+    Query,
+    QueryCreate,
+    QueryResults,
+    QueryState,
+    StatementResults,
+    decode_results,
+    encode_results,
+)
+from djqs.utils import get_settings
+
+_logger = logging.getLogger(__name__)
+router = APIRouter(tags=["SQL Queries"])
+
+
+@router.post(
+    "/queries/",
+    response_model=QueryResults,
+    status_code=HTTPStatus.OK,
+    responses={
+        200: {
+            "content": {"application/msgpack": {}},
+            "description": "Return results as JSON or msgpack",
+        },
+    },
+)
+async def submit_query(  # pylint: disable=too-many-arguments
+    accept: Optional[str] = Header(None),
+    *,
+    settings: Settings = Depends(get_settings),
+    request: Request,
+    response: Response,
+    postgres_pool: AsyncConnectionPool = Depends(get_postgres_pool),
+    background_tasks: BackgroundTasks,
+    body: Any = Body(
+        ...,
+        example={
+            "catalog_name": "warehouse",
+            "engine_name": "trino",
+            "engine_version": "451",
+            "submitted_query": "select * from tpch.sf1.customer limit 10",
+        },
+    ),
+) -> QueryResults:
+    """
+    Run or schedule a query.
+
+    This endpoint is different from others in that it accepts both JSON and msgpack, and
+    can also return JSON or msgpack, depending on HTTP headers.
+    """
+    content_type = request.headers.get("content-type")
+    if content_type == "application/json":
+        data = body
+    elif content_type == "application/msgpack":
+        data = json.loads(msgpack.unpackb(body, ext_hook=decode_results))
+    elif content_type is None:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST,
+            detail="Content type must be specified",
+        )
+    else:
+        raise HTTPException(
+            status_code=HTTPStatus.UNPROCESSABLE_ENTITY,
+            detail=f"Content type not accepted: {content_type}",
+        )
+
+    # Set default catalog and engine if not explicitly specified in submitted query
+    data["engine_name"] = data.get("engine_name") or settings.default_engine
+    data["engine_version"] = (
+        data.get("engine_version") or settings.default_engine_version
+    )
+    data["catalog_name"] = data.get("catalog_name") or settings.default_catalog
+
+    create_query = QueryCreate(**data)
+
+    query_with_results = await save_query_and_run(
+        create_query=create_query,
+        settings=settings,
+        response=response,
+        background_tasks=background_tasks,
+        postgres_pool=postgres_pool,
+        headers=request.headers,
+    )
+
+    return_type = get_best_match(accept, ["application/json", "application/msgpack"])
+    if not return_type:
+        raise HTTPException(
+            status_code=HTTPStatus.NOT_ACCEPTABLE,
+            detail="Client MUST accept: application/json, application/msgpack",
+        )
+
+    if return_type == "application/msgpack":
+        content = msgpack.packb(
+            asdict(query_with_results),
+            default=encode_results,
+        )
+    else:
+        content = json.dumps(asdict(query_with_results), default=str)
+
+    return Response(
+        content=content,
+        media_type=return_type,
+        status_code=response.status_code or HTTPStatus.OK,
+    )
+
+
+async def save_query_and_run(  # pylint: disable=R0913
+    create_query: QueryCreate,
+    settings: Settings,
+    response: Response,
+    background_tasks: BackgroundTasks,
+    postgres_pool: AsyncConnectionPool,
+    headers: Optional[Dict[str, str]] = None,
+) -> QueryResults:
+    """
+    Store a new query to the DB and run it.
+    """
+    query = Query(
+        catalog_name=create_query.catalog_name,  # type: ignore
+        engine_name=create_query.engine_name,  # type: ignore
+        engine_version=create_query.engine_version,  # type: ignore
+        submitted_query=create_query.submitted_query,
+        async_=create_query.async_,
+    )
+    query.state = QueryState.ACCEPTED
+
+    async with postgres_pool.connection() as conn:
+        results = (
+            await DBQuery()
+            .save_query(
+                query_id=query.id,
+                catalog_name=query.catalog_name,
+                engine_name=query.engine_name,
+                engine_version=query.engine_version,
+                submitted_query=query.submitted_query,
+                async_=query.async_,
+                state=query.state.value,
+            )
+            .execute(conn=conn)
+        )
+        query_save_result = results[0]
+        if not query_save_result:  # pragma: no cover
+            raise HTTPException(
+                status_code=HTTPStatus.NOT_FOUND,
+                detail="Query failed to save",
+            )
+
+        if query.async_:
+            background_tasks.add_task(
+                process_query,
+                settings,
+                postgres_pool,
+                query,
+                headers,
+            )
+
+            response.status_code = HTTPStatus.CREATED
+            return QueryResults(
+                id=query.id,
+                catalog_name=query.catalog_name,
+                engine_name=query.engine_name,
+                engine_version=query.engine_version,
+                submitted_query=query.submitted_query,
+                executed_query=query.executed_query,
+                state=QueryState.SCHEDULED,
+                results=[],
+                errors=[],
+            )
+
+    query_results = await process_query(
+        settings=settings,
+        postgres_pool=postgres_pool,
+        query=query,
+        headers=headers,
+    )
+    return query_results
+
+
+def load_query_results(
+    settings: Settings,
+    key: str,
+) -> List[StatementResults]:
+    """
+    Load results from backend, if available.
+
+    If ``paginate`` is true we also load the results into the cache, anticipating more
+    paginated queries.
+    """
+    if settings.results_backend.has(key):
+        _logger.info("Reading results from results backend")
+        cached = settings.results_backend.get(key)
+        query_results = json.loads(cached)
+    else:  # pragma: no cover
+        _logger.warning("No results found")
+        query_results = []
+
+    return query_results
+
+
+@router.get("/queries/{query_id}/", response_model=QueryResults)
+async def read_query(
+    query_id: uuid.UUID,
+    *,
+    settings: Settings = Depends(get_settings),
+    postgres_pool: AsyncConnectionPool = Depends(get_postgres_pool),
+) -> QueryResults:
+    """
+    Fetch information about a query.
+
+    For paginated queries we move the data from the results backend to the cache for a
+    short period, anticipating additional requests.
+    """
+    async with postgres_pool.connection() as conn:
+        dbquery_results = (
+            await DBQuery().get_query(query_id=query_id).execute(conn=conn)
+        )
+        queries = dbquery_results[0]
+        if not queries:
+            raise HTTPException(
+                status_code=HTTPStatus.NOT_FOUND,
+                detail="Query not found",
+            )
+        query = queries[0]
+
+    query_results = load_query_results(settings, str(query_id))
+
+    prev = next_ = None
+
+    return QueryResults(
+        results=query_results, next=next_, previous=prev, errors=[], **query
+    )

{datajunction_query-0.0.1a65 → datajunction_query-0.0.1a66}/djqs/api/tables.py

@@ -3,24 +3,21 @@ Table related APIs.
 """
 from typing import Optional
 
-from fastapi import APIRouter, Depends
-from sqlmodel import Session
+from fastapi import APIRouter, Path, Query
 
-from djqs.api.helpers import get_columns, get_engine
+from djqs.api.helpers import get_columns
 from djqs.exceptions import DJInvalidTableRef
 from djqs.models.table import TableInfo
-from djqs.utils import get_session, get_settings
+from djqs.utils import get_settings
 
 router = APIRouter(tags=["Table Reflection"])
 
 
 @router.get("/table/{table}/columns/", response_model=TableInfo)
 def table_columns(
-    table: str,
-    engine: Optional[str] = None,
-    engine_version: Optional[str] = None,
-    *,
-    session: Session = Depends(get_session),
+    table: str = Path(..., example="tpch.sf1.customer"),
+    engine: Optional[str] = Query(None, example="trino"),
+    engine_version: Optional[str] = Query(None, example="451"),
 ) -> TableInfo:
     """
     Get column information for a table
@@ -34,19 +31,19 @@ def table_columns(
         )
     settings = get_settings()
 
-    if engine_version == "":
-        version = ""
-    else:  # pragma: no cover
-        version = engine_version or settings.default_reflection_engine_version
-
-    engine = get_engine(
-        session=session,
-        name=engine or settings.default_reflection_engine,
-        version=version,
-    )
+    if engine and engine_version:
+        engine_config = settings.find_engine(
+            engine_name=engine,
+            engine_version=engine_version or settings.default_engine_version,
+        )
+    else:
+        engine_config = settings.find_engine(
+            engine_name=settings.default_engine,
+            engine_version=engine_version or settings.default_engine_version,
+        )
     external_columns = get_columns(
-        uri=engine.uri,
-        extra_params=engine.extra_params,
+        uri=engine_config.uri,
+        extra_params=engine_config.extra_params,
         catalog=table_parts[0],
         schema=table_parts[1],
         table=table_parts[2],