datajunction-query 0.0.1a61__tar.gz → 0.0.1a62__tar.gz

{datajunction_query-0.0.1a61 → datajunction_query-0.0.1a62}/.pre-commit-config.yaml

@@ -69,7 +69,7 @@ repos:
     - types-requests
     - types-freezegun
     - types-python-dateutil
-    - types-pkg_resources
+    - types-setuptools
     - types-PyYAML
     - types-tabulate
     - types-toml

{datajunction_query-0.0.1a61 → datajunction_query-0.0.1a62}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: datajunction-query
-Version: 0.0.1a61
+Version: 0.0.1a62
 Summary: OSS Implementation of a DataJunction Query Service
 Project-URL: repository, https://github.com/DataJunction/dj
 Author-email: DataJunction Authors <roberto@dealmeida.net>

{datajunction_query-0.0.1a61 → datajunction_query-0.0.1a62}/djqs/__about__.py

@@ -1,4 +1,4 @@
 """
 Version for Hatch
 """
-__version__ = "0.0.1a61"
+__version__ = "0.0.1a62"

{datajunction_query-0.0.1a61 → datajunction_query-0.0.1a62}/djqs/api/main.py

@@ -32,7 +32,7 @@ async def lifespan(app: FastAPI):  # pylint: disable=W0621,W0613
     """
     try:
         load_djqs_config(settings=settings, session=session)
-    except Exception as e:  # pylint: disable=W0718
+    except Exception as e:  # pylint: disable=W0718,C0103
         _logger.warning("Could not load DJQS config: %s", e)
     yield
 

{datajunction_query-0.0.1a61 → datajunction_query-0.0.1a62}/djqs/api/queries.py

@@ -5,7 +5,7 @@ import json
 import logging
 import uuid
 from http import HTTPStatus
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 
 import msgpack
 from accept_types import get_best_match
@@ -105,6 +105,7 @@ async def submit_query(  # pylint: disable=too-many-arguments
         settings,
         response,
         background_tasks,
+        request.headers,
     )
 
     return_type = get_best_match(accept, ["application/json", "application/msgpack"])
@@ -129,12 +130,13 @@ async def submit_query(  # pylint: disable=too-many-arguments
     )
 
 
-def save_query_and_run(
+def save_query_and_run(  # pylint: disable=R0913
     create_query: QueryCreate,
     session: Session,
     settings: Settings,
     response: Response,
     background_tasks: BackgroundTasks,
+    headers: Optional[Dict[str, str]] = None,
 ) -> QueryResults:
     """
     Store a new query to the DB and run it.
@@ -147,12 +149,12 @@ def save_query_and_run(
     session.refresh(query)
 
     if query.async_:
-        background_tasks.add_task(process_query, session, settings, query)
+        background_tasks.add_task(process_query, session, settings, query, headers)
 
         response.status_code = HTTPStatus.CREATED
         return QueryResults(results=[], errors=[], **query.dict())
 
-    return process_query(session, settings, query)
+    return process_query(session, settings, query, headers)
 
 
 def load_query_results(

{datajunction_query-0.0.1a61 → datajunction_query-0.0.1a62}/djqs/constants.py

@@ -15,3 +15,6 @@ DEFAULT_DIMENSION_COLUMN = "id"
 # used by the SQLAlchemy client
 QUERY_EXECUTE_TIMEOUT = timedelta(seconds=60)
 GET_COLUMNS_TIMEOUT = timedelta(seconds=60)
+
+# Request header configuration params
+SQLALCHEMY_URI = "SQLALCHEMY_URI"

{datajunction_query-0.0.1a61 → datajunction_query-0.0.1a62}/djqs/engine.py

@@ -1,11 +1,10 @@
 """
 Query related functions.
 """
-
 import logging
 import os
 from datetime import datetime, timezone
-from typing import List, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import duckdb
 import snowflake.connector
@@ -14,6 +13,7 @@ from sqlalchemy import create_engine, text
 from sqlmodel import Session, select
 
 from djqs.config import Settings
+from djqs.constants import SQLALCHEMY_URI
 from djqs.models.engine import Engine, EngineType
 from djqs.models.query import (
     ColumnMetadata,
@@ -66,9 +66,10 @@ def get_columns_from_description(
     return columns
 
 
-def run_query(
+def run_query(  # pylint: disable=R0914
     session: Session,
     query: Query,
+    headers: Optional[Dict[str, str]] = None,
 ) -> List[Tuple[str, List[ColumnMetadata], Stream]]:
     """
     Run a query and return its results.
@@ -76,13 +77,25 @@ def run_query(
     For each statement we return a tuple with the statement SQL, a description of the
     columns (name and type) and a stream of rows (tuples).
     """
+
     _logger.info("Running query on catalog %s", query.catalog_name)
+
     engine = session.exec(
         select(Engine)
         .where(Engine.name == query.engine_name)
         .where(Engine.version == query.engine_version),
     ).one()
-    if engine.type == EngineType.DUCKDB:
+
+    query_server = headers.get("SQLALCHEMY_URI") if headers else None
+
+    if query_server:
+        _logger.info(
+            "Creating sqlalchemy engine using request header param %s",
+            SQLALCHEMY_URI,
+        )
+        sqla_engine = create_engine(query_server)
+    elif engine.type == EngineType.DUCKDB:
+        _logger.info("Creating duckdb connection")
         conn = (
             duckdb.connect()
             if engine.uri == "duckdb:///:memory:"
@@ -92,7 +105,8 @@ def run_query(
             )
         )
         return run_duckdb_query(query, conn)
-    if engine.type == EngineType.SNOWFLAKE:
+    elif engine.type == EngineType.SNOWFLAKE:
+        _logger.info("Creating snowflake connection")
         conn = snowflake.connector.connect(
             **engine.extra_params,
             password=os.getenv("SNOWSQL_PWD"),
@@ -100,6 +114,10 @@ def run_query(
         cur = conn.cursor()
 
         return run_snowflake_query(query, cur)
+
+    _logger.info(
+        "Creating sqlalchemy engine using engine name and version defined on query",
+    )
     sqla_engine = create_engine(engine.uri, connect_args=engine.extra_params)
     connection = sqla_engine.connect()
 
@@ -152,6 +170,7 @@ def process_query(
     session: Session,
     settings: Settings,
     query: Query,
+    headers: Optional[Dict[str, str]] = None,
 ) -> QueryResults:
     """
     Process a query.
@@ -164,7 +183,11 @@ def process_query(
     query.started = datetime.now(timezone.utc)
     try:
         root = []
-        for sql, columns, stream in run_query(session=session, query=query):
+        for sql, columns, stream in run_query(
+            session=session,
+            query=query,
+            headers=headers,
+        ):
             rows = list(stream)
             root.append(
                 StatementResults(

{datajunction_query-0.0.1a61 → datajunction_query-0.0.1a62}/tests/api/queries_test.py

@@ -86,6 +86,71 @@ def test_submit_query(session: Session, client: TestClient) -> None:
     assert data["errors"] == []
 
 
+def test_submit_query_with_sqlalchemy_uri_header(
+    session: Session,
+    client: TestClient,
+) -> None:
+    """
+    Test ``POST /queries/`` with the SQLALCHEMY_URI defined in the header.
+    """
+    engine = Engine(
+        name="test_engine",
+        type=EngineType.DUCKDB,
+        version="1.0",
+        uri="duckdb:///:memory:",
+    )
+    catalog = Catalog(name="test_catalog", engines=[engine])
+    session.add(catalog)
+    session.commit()
+    session.refresh(catalog)
+
+    query_create = QueryCreate(
+        catalog_name=catalog.name,
+        engine_name=engine.name,
+        engine_version=engine.version,
+        submitted_query="SELECT 1 AS col",
+    )
+    payload = query_create.json(by_alias=True)
+    assert payload == json.dumps(
+        {
+            "catalog_name": "test_catalog",
+            "engine_name": "test_engine",
+            "engine_version": "1.0",
+            "submitted_query": "SELECT 1 AS col",
+            "async_": False,
+        },
+    )
+
+    with freeze_time("2021-01-01T00:00:00Z"):
+        response = client.post(
+            "/queries/",
+            data=payload,
+            headers={
+                "Content-Type": "application/json",
+                "Accept": "application/json",
+                "SQLALCHEMY_URI": "trino://example@foo.bar/catalog/schema",
+            },
+        )
+    data = response.json()
+
+    assert response.status_code == 200
+    assert data["catalog_name"] == "test_catalog"
+    assert data["engine_name"] == "test_engine"
+    assert data["engine_version"] == "1.0"
+    assert data["submitted_query"] == "SELECT 1 AS col"
+    assert data["executed_query"] == "SELECT 1 AS col"
+    assert data["scheduled"] == "2021-01-01T00:00:00"
+    assert data["started"] == "2021-01-01T00:00:00"
+    assert data["finished"] == "2021-01-01T00:00:00"
+    assert data["state"] == "FINISHED"
+    assert data["progress"] == 1.0
+    assert len(data["results"]) == 1
+    assert data["results"][0]["sql"] == "SELECT 1 AS col"
+    assert data["results"][0]["columns"] == [{"name": "col", "type": "STR"}]
+    assert data["results"][0]["rows"] == [[1]]
+    assert data["errors"] == []
+
+
 def test_submit_query_msgpack(session: Session, client: TestClient) -> None:
     """
     Test ``POST /queries/`` using msgpack.