datajunction-query 0.0.1a64__py3-none-any.whl → 0.0.1a66__py3-none-any.whl

{datajunction_query-0.0.1a64.dist-info → datajunction_query-0.0.1a66.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: datajunction-query
-Version: 0.0.1a64
+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.1a66.dist-info/RECORD

@@ -0,0 +1,24 @@
+djqs/__about__.py,sha256=mKc8a1KL2DOwHh4C4Sfmqi_Cpqvps_-YclpQjIrRYLM,51
+djqs/__init__.py,sha256=nN5-uJoSVEwuc8n-wMygqeF0Xhxi_zqqbCgutZvAt3E,384
+djqs/config.py,sha256=8WLfe3M3maO4WhWb4ukrj8gdIlHQyLC4j9Q7KQ95bOQ,6662
+djqs/constants.py,sha256=WS0uC-cKLQ_ZPhtevcv6tJ8WI6UtxYGZwTqJJtsTep0,485
+djqs/engine.py,sha256=o_WhKhYJ4Ld6RB9KuRb_J7exZVo3yQ9ztId_FbOBzIQ,7340
+djqs/enum.py,sha256=GJVLYDJ2zWjjUBENgyjZZ_94A24BJtvci-AKTmA0zek,590
+djqs/exceptions.py,sha256=7S5hU6i9umAtITIDCq5_NIC8obTnF_f8Z7xTutyQU1I,5872
+djqs/fixes.py,sha256=TcXnh0I1z4vEAupPatzrvnqyToGihndnxmLJtIn-_Z8,33
+djqs/typing.py,sha256=TpZHhrK_lzEYg_ZlT5qVCJz8seQBKwrULiTPO-lMxEU,6220
+djqs/utils.py,sha256=iTctVXY4uSQnGn1O-_CdDXldiOwfufz582Lo8izox9M,906
+djqs/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+djqs/api/helpers.py,sha256=VjMIIz6Se9GUCE8xfAEtNfPqjCWNRburjjX3AWTcAMY,1358
+djqs/api/main.py,sha256=SR0vbX0W1v8orfluyfsALrT1G3bF9emH1KBWoO9lcyw,2237
+djqs/api/queries.py,sha256=xh11sIbFlKcnOHfbOLCf4ZWzAAhJ4m4FVq0oYMSPRO8,7709
+djqs/api/tables.py,sha256=oxjSMDUxtuCzvXP58Oyo0Ji_T0BnNrpGo9wQ1hv9zeQ,1674
+djqs/db/postgres.py,sha256=ggTAVjm52oRbdfgHg1e0QXXb7AkYgFP_bGBys7IP5yE,4862
+djqs/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+djqs/models/query.py,sha256=n02TPSq8cZOuH49FyVieVwhGdrkaZ-4NpqCSuOgL5X4,3324
+djqs/models/table.py,sha256=7HmiXWHBWEthqdwPYLlsOgeKD-w9urCHoC4xaQ1RnM8,238
+datajunction_query-0.0.1a66.dist-info/METADATA,sha256=8MIrqCyR5fSCey6lArtDwud-OWccAsc1r-Vtcz6-tJY,9525
+datajunction_query-0.0.1a66.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+datajunction_query-0.0.1a66.dist-info/licenses/AUTHORS.rst,sha256=G9YmXPfQ0kAuxqlCwvWVvUnZitP9hAc-rPfZ5q7Pv1U,235
+datajunction_query-0.0.1a66.dist-info/licenses/LICENSE.txt,sha256=KuSxhVgPuUGoYWphJig4POcTAIUNLUj8vOx-cqQFMj8,1081
+datajunction_query-0.0.1a66.dist-info/RECORD,,

djqs/__about__.py

@@ -1,4 +1,4 @@
 """
 Version for Hatch
 """
-__version__ = "0.0.1a64"
+__version__ = "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(

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)

djqs/api/queries.py

@@ -1,9 +1,11 @@
 """
 Query related APIs.
 """
+
 import json
 import logging
 import uuid
+from dataclasses import asdict
 from http import HTTPStatus
 from typing import Any, Dict, List, Optional
 
@@ -19,21 +21,21 @@ from fastapi import (
     Request,
     Response,
 )
-from sqlmodel import Session
+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,
-    Results,
     StatementResults,
     decode_results,
     encode_results,
 )
-from djqs.utils import get_session, get_settings
+from djqs.utils import get_settings
 
 _logger = logging.getLogger(__name__)
 router = APIRouter(tags=["SQL Queries"])
@@ -49,32 +51,24 @@ router = APIRouter(tags=["SQL Queries"])
             "description": "Return results as JSON or msgpack",
         },
     },
-    openapi_extra={
-        "requestBody": {
-            "content": {
-                "application/json": {
-                    "schema": QueryCreate.schema(
-                        ref_template="#/components/schemas/{model}",
-                    ),
-                },
-                "application/msgpack": {
-                    "schema": QueryCreate.schema(
-                        ref_template="#/components/schemas/{model}",
-                    ),
-                },
-            },
-        },
-    },
 )
 async def submit_query(  # pylint: disable=too-many-arguments
     accept: Optional[str] = Header(None),
     *,
-    session: Session = Depends(get_session),
     settings: Settings = Depends(get_settings),
     request: Request,
     response: Response,
+    postgres_pool: AsyncConnectionPool = Depends(get_postgres_pool),
     background_tasks: BackgroundTasks,
-    body: Any = Body(...),
+    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.
@@ -86,7 +80,7 @@ async def submit_query(  # pylint: disable=too-many-arguments
     if content_type == "application/json":
         data = body
     elif content_type == "application/msgpack":
-        data = msgpack.unpackb(body, ext_hook=decode_results)
+        data = json.loads(msgpack.unpackb(body, ext_hook=decode_results))
     elif content_type is None:
         raise HTTPException(
             status_code=HTTPStatus.BAD_REQUEST,
@@ -97,15 +91,23 @@ async def submit_query(  # pylint: disable=too-many-arguments
             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 = save_query_and_run(
-        create_query,
-        session,
-        settings,
-        response,
-        background_tasks,
-        request.headers,
+    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"])
@@ -117,11 +119,11 @@ async def submit_query(  # pylint: disable=too-many-arguments
 
     if return_type == "application/msgpack":
         content = msgpack.packb(
-            query_with_results.dict(by_alias=True),
+            asdict(query_with_results),
             default=encode_results,
         )
     else:
-        content = query_with_results.json(by_alias=True)
+        content = json.dumps(asdict(query_with_results), default=str)
 
     return Response(
         content=content,
@@ -130,31 +132,76 @@ async def submit_query(  # pylint: disable=too-many-arguments
     )
 
 
-def save_query_and_run(  # pylint: disable=R0913
+async def save_query_and_run(  # pylint: disable=R0913
     create_query: QueryCreate,
-    session: Session,
     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(**create_query.dict(by_alias=True))
+    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
 
-    session.add(query)
-    session.commit()
-    session.refresh(query)
+    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, session, settings, query, headers)
+        if query.async_:
+            background_tasks.add_task(
+                process_query,
+                settings,
+                postgres_pool,
+                query,
+                headers,
+            )
 
-        response.status_code = HTTPStatus.CREATED
-        return QueryResults(results=[], errors=[], **query.dict())
+            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=[],
+            )
 
-    return process_query(session, settings, query, headers)
+    query_results = await process_query(
+        settings=settings,
+        postgres_pool=postgres_pool,
+        query=query,
+        headers=headers,
+    )
+    return query_results
 
 
 def load_query_results(
@@ -171,7 +218,7 @@ def load_query_results(
         _logger.info("Reading results from results backend")
         cached = settings.results_backend.get(key)
         query_results = json.loads(cached)
-    else:
+    else:  # pragma: no cover
         _logger.warning("No results found")
         query_results = []
 
@@ -179,11 +226,11 @@ def load_query_results(
 
 
 @router.get("/queries/{query_id}/", response_model=QueryResults)
-def read_query(
+async def read_query(
     query_id: uuid.UUID,
     *,
-    session: Session = Depends(get_session),
     settings: Settings = Depends(get_settings),
+    postgres_pool: AsyncConnectionPool = Depends(get_postgres_pool),
 ) -> QueryResults:
     """
     Fetch information about a query.
@@ -191,15 +238,22 @@ def read_query(
     For paginated queries we move the data from the results backend to the cache for a
     short period, anticipating additional requests.
     """
-    query = session.get(Query, query_id)
-    if not query:
-        raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Query not found")
+    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
-    results = Results(__root__=query_results)
 
     return QueryResults(
-        results=results, next=next_, previous=prev, errors=[], **query.dict()
+        results=query_results, next=next_, previous=prev, errors=[], **query
     )

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],