datajunction-query 0.0.1a61__py3-none-any.whl → 0.0.1a62__py3-none-any.whl

{datajunction_query-0.0.1a61.dist-info → datajunction_query-0.0.1a62.dist-info}/METADATA

@@ -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.dist-info → datajunction_query-0.0.1a62.dist-info}/RECORD

@@ -1,8 +1,8 @@
-djqs/__about__.py,sha256=YIom-Iew--PmtjAWNmL7cutQitEoQOyCSRG0RfhsLwc,51
+djqs/__about__.py,sha256=CutNphqU7kSQIFVLpZyS_P1sGJgbzf_o5Rh_gCBvt2E,51
 djqs/__init__.py,sha256=nN5-uJoSVEwuc8n-wMygqeF0Xhxi_zqqbCgutZvAt3E,384
 djqs/config.py,sha256=Flva9W8McnTunUlHGhYd0CsgbsjWbZIoPW25ZTzFxws,3029
-djqs/constants.py,sha256=T-sO5C5-NNdEywGglTIu6bjT1VooFwW4XWNA_ZwSYM8,412
-djqs/engine.py,sha256=5HvrFnQmy89j2duPNT-zSktpfl7zKj-FZdeJ5SvFf8M,5719
+djqs/constants.py,sha256=WS0uC-cKLQ_ZPhtevcv6tJ8WI6UtxYGZwTqJJtsTep0,485
+djqs/engine.py,sha256=H3tSKVgBBc24OWMwl6Acz1DjaciF23kBik98w2cVXUU,6453
 djqs/enum.py,sha256=GJVLYDJ2zWjjUBENgyjZZ_94A24BJtvci-AKTmA0zek,590
 djqs/exceptions.py,sha256=z-EjGLRwHL_v3zl1-RICwQis7wUwl3000tse0YwuoBE,4992
 djqs/fixes.py,sha256=TcXnh0I1z4vEAupPatzrvnqyToGihndnxmLJtIn-_Z8,33
@@ -12,16 +12,16 @@ djqs/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 djqs/api/catalogs.py,sha256=I_i8-Nok0V3ND1FpOCQ5Nmgaof5MLH8i0oeTjvDxQgw,3062
 djqs/api/engines.py,sha256=278fsgIgQXe4WAmbfbxoA08ALe57us_zKlaK_hbiDrc,1611
 djqs/api/helpers.py,sha256=91UGqs-SAq3fCU9yElBMev9vqtDospOjBq0icfFl1qs,2468
-djqs/api/main.py,sha256=cbZHxkfmushgw8e0G2q8f5Ytexqk-qjKtZPyHK5sbjQ,2075
-djqs/api/queries.py,sha256=gSdzI7pVhCsZDpjP5l0NTvHq4ykF20CFzKiLZrrlUmQ,5523
+djqs/api/main.py,sha256=Sjhh41RIoj-d1BQpBquW9D8O0n0Ft9Ka2Ntq7i_lo3o,2081
+djqs/api/queries.py,sha256=eCs7dKt0mOJ8tGUNIk0nomWG7m_VPsnjUUJ7c7_dvuA,5643
 djqs/api/tables.py,sha256=1hwYPRAG9kZfS8HE3r_omC4AjHadAbdhFIwx3e4KOvo,1582
 djqs/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 djqs/models/catalog.py,sha256=o3Wlghwr6ax4Eo-3zE5Mdq478xyUxeaf2o3whFprXtU,2038
 djqs/models/engine.py,sha256=k9cBvnYAfzz4yaBDvErDqXIwlqrRsKRbGikBKTQHJPA,889
 djqs/models/query.py,sha256=Qsy69Cfrkn5myOZZfLrR83JB4ynPXogitUo8A7hRsz8,3338
 djqs/models/table.py,sha256=7HmiXWHBWEthqdwPYLlsOgeKD-w9urCHoC4xaQ1RnM8,238
-datajunction_query-0.0.1a61.dist-info/METADATA,sha256=AVYv4rZeQ91zjRdChyV50YIhTpD2FCVO78iJ_LbAcKw,9495
-datajunction_query-0.0.1a61.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-datajunction_query-0.0.1a61.dist-info/licenses/AUTHORS.rst,sha256=G9YmXPfQ0kAuxqlCwvWVvUnZitP9hAc-rPfZ5q7Pv1U,235
-datajunction_query-0.0.1a61.dist-info/licenses/LICENSE.txt,sha256=KuSxhVgPuUGoYWphJig4POcTAIUNLUj8vOx-cqQFMj8,1081
-datajunction_query-0.0.1a61.dist-info/RECORD,,
+datajunction_query-0.0.1a62.dist-info/METADATA,sha256=3L3Cg3n-OUE9vopB_3GNzJKNppk0_aoSLJz4-St5KOs,9495
+datajunction_query-0.0.1a62.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+datajunction_query-0.0.1a62.dist-info/licenses/AUTHORS.rst,sha256=G9YmXPfQ0kAuxqlCwvWVvUnZitP9hAc-rPfZ5q7Pv1U,235
+datajunction_query-0.0.1a62.dist-info/licenses/LICENSE.txt,sha256=KuSxhVgPuUGoYWphJig4POcTAIUNLUj8vOx-cqQFMj8,1081
+datajunction_query-0.0.1a62.dist-info/RECORD,,

djqs/__about__.py

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

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(

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"

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(