datajunction-query 0.0.1a58__py3-none-any.whl → 0.0.28__py3-none-any.whl

djqs/config.py

@@ -3,94 +3,197 @@ Configuration for the query service
 """
 
 import json
+import logging
+import os
 from datetime import timedelta
-from typing import Optional
+from enum import Enum
+from typing import Dict, List, Optional
 
 import toml
 import yaml
 from cachelib.base import BaseCache
 from cachelib.file import FileSystemCache
-from pydantic import BaseSettings
-from sqlmodel import Session, delete, select
 
-from djqs.exceptions import DJException
-from djqs.models.catalog import Catalog, CatalogEngines
-from djqs.models.engine import Engine
+from djqs.exceptions import DJUnknownCatalog, DJUnknownEngine
 
+_logger = logging.getLogger(__name__)
 
-class Settings(BaseSettings):  # pylint: disable=too-few-public-methods
+
+class EngineType(Enum):
     """
-    Configuration for the query service
+    Supported engine types
     """
 
-    name: str = "DJQS"
-    description: str = "A DataJunction Query Service"
-    url: str = "http://localhost:8001/"
+    DUCKDB = "duckdb"
+    SQLALCHEMY = "sqlalchemy"
+    SNOWFLAKE = "snowflake"
+    TRINO = "trino"
 
-    # SQLAlchemy URI for the metadata database.
-    index: str = "sqlite:///djqs.db?check_same_thread=False"
 
-    # The default engine to use for reflection
-    default_reflection_engine: str = "default"
+class EngineInfo:  # pylint: disable=too-few-public-methods
+    """
+    Information about a query engine
+    """
 
-    # The default engine version to use for reflection
-    default_reflection_engine_version: str = ""
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        name: str,
+        version: str,
+        type: str,  # pylint: disable=redefined-builtin
+        uri: str,
+        extra_params: Optional[Dict[str, str]] = None,
+    ):
+        self.name = name
+        self.version = str(version)
+        self.type = EngineType(type)
+        self.uri = uri
+        self.extra_params = extra_params or {}
+
+
+class CatalogInfo:  # pylint: disable=too-few-public-methods
+    """
+    Information about a catalog
+    """
 
-    # Where to store the results from queries.
-    results_backend: BaseCache = FileSystemCache("/tmp/djqs", default_timeout=0)
+    def __init__(self, name: str, engines: List[str]):
+        self.name = name
+        self.engines = engines
 
-    paginating_timeout: timedelta = timedelta(minutes=5)
 
-    # How long to wait when pinging databases to find out the fastest online database.
-    do_ping_timeout: timedelta = timedelta(seconds=5)
+class Settings:  # pylint: disable=too-many-instance-attributes
+    """
+    Configuration for the query service
+    """
 
-    # Configuration file for catalogs and engines
-    configuration_file: Optional[str] = None
+    def __init__(  # pylint: disable=too-many-arguments,too-many-locals,dangerous-default-value
+        self,
+        name: Optional[str] = "DJQS",
+        description: Optional[str] = "A DataJunction Query Service",
+        url: Optional[str] = "http://localhost:8001/",
+        index: Optional[str] = "postgresql://dj:dj@postgres_metadata:5432/djqs",
+        default_catalog: Optional[str] = "",
+        default_engine: Optional[str] = "",
+        default_engine_version: Optional[str] = "",
+        results_backend: Optional[BaseCache] = None,
+        results_backend_path: Optional[str] = "/tmp/djqs",
+        results_backend_timeout: Optional[str] = "0",
+        paginating_timeout_minutes: Optional[str] = "5",
+        do_ping_timeout_seconds: Optional[str] = "5",
+        configuration_file: Optional[str] = None,
+        engines: Optional[List[EngineInfo]] = None,
+        catalogs: Optional[List[CatalogInfo]] = None,
+    ):
+        self.name: str = os.getenv("NAME", name or "")
+        self.description: str = os.getenv("DESCRIPTION", description or "")
+        self.url: str = os.getenv("URL", url or "")
+
+        # SQLAlchemy URI for the metadata database.
+        self.index: str = os.getenv("INDEX", index or "")
+
+        # The default catalog to use if not specified in query payload
+        self.default_catalog: str = os.getenv("DEFAULT_CATALOG", default_catalog or "")
+
+        # The default engine to use if not specified in query payload
+        self.default_engine: str = os.getenv("DEFAULT_ENGINE", default_engine or "")
+
+        # The default engine version to use if not specified in query payload
+        self.default_engine_version: str = os.getenv(
+            "DEFAULT_ENGINE_VERSION",
+            default_engine_version or "",
+        )
 
-    # Enable setting catalog and engine config via REST API calls
-    enable_dynamic_config: bool = True
+        # Where to store the results from queries.
+        self.results_backend: BaseCache = results_backend or FileSystemCache(
+            os.getenv("RESULTS_BACKEND_PATH", results_backend_path or ""),
+            default_timeout=int(
+                os.getenv("RESULTS_BACKEND_TIMEOUT", results_backend_timeout or "0"),
+            ),
+        )
 
+        self.paginating_timeout: timedelta = timedelta(
+            minutes=int(
+                os.getenv(
+                    "PAGINATING_TIMEOUT_MINUTES",
+                    paginating_timeout_minutes or "5",
+                ),
+            ),
+        )
 
-def load_djqs_config(settings: Settings, session: Session) -> None:  # pragma: no cover
-    """
-    Load the DJQS config file into the server metadata database
-    """
-    config_file = settings.configuration_file if settings.configuration_file else None
-    if not config_file:
-        return
-
-    session.exec(delete(Catalog))
-    session.exec(delete(Engine))
-    session.exec(delete(CatalogEngines))
-    session.commit()
-
-    with open(config_file, mode="r", encoding="utf-8") as filestream:
-
-        def unknown_filetype():
-            raise DJException(message=f"Unknown config file type: {config_file}")
-
-        data = (
-            yaml.safe_load(filestream)
-            if any([config_file.endswith(".yml"), config_file.endswith(".yaml")])
-            else toml.load(filestream)
-            if config_file.endswith(".toml")
-            else json.load(filestream)
-            if config_file.endswith(".json")
-            else unknown_filetype()
+        # How long to wait when pinging databases to find out the fastest online database.
+        self.do_ping_timeout: timedelta = timedelta(
+            seconds=int(
+                os.getenv("DO_PING_TIMEOUT_SECONDS", do_ping_timeout_seconds or "5"),
+            ),
         )
 
-    for engine in data["engines"]:
-        session.add(Engine.parse_obj(engine))
-    session.commit()
+        # Configuration file for catalogs and engines
+        self.configuration_file: Optional[str] = (
+            os.getenv("CONFIGURATION_FILE") or configuration_file
+        )
 
-    for catalog in data["catalogs"]:
-        attached_engines = []
-        catalog_engines = catalog.pop("engines")
-        for name in catalog_engines:
-            attached_engines.append(
-                session.exec(select(Engine).where(Engine.name == name)).one(),
+        self.engines: List[EngineInfo] = engines or []
+        self.catalogs: List[CatalogInfo] = catalogs or []
+
+        self._load_configuration()
+
+    def _load_configuration(self):
+        config_file = self.configuration_file
+
+        if config_file:
+            if config_file.endswith(".yaml") or config_file.endswith(".yml"):
+                with open(config_file, "r", encoding="utf-8") as file:
+                    config = yaml.safe_load(file)
+            elif config_file.endswith(".toml"):
+                with open(config_file, "r", encoding="utf-8") as file:
+                    config = toml.load(file)
+            elif config_file.endswith(".json"):
+                with open(config_file, "r", encoding="utf-8") as file:
+                    config = json.load(file)
+            else:
+                raise ValueError(
+                    f"Unsupported configuration file format: {config_file}",
+                )
+
+            self.engines = [
+                EngineInfo(**engine) for engine in config.get("engines", [])
+            ]
+            self.catalogs = [
+                CatalogInfo(**catalog) for catalog in config.get("catalogs", [])
+            ]
+        else:
+            _logger.warning("No settings configuration file has been set")
+
+    def find_engine(
+        self,
+        engine_name: str,
+        engine_version: str,
+    ) -> EngineInfo:
+        """
+        Find an engine defined in the server configuration
+        """
+        found_engine = None
+        for engine in self.engines:
+            if engine.name == engine_name and engine.version == engine_version:
+                found_engine = engine
+        if not found_engine:
+            raise DJUnknownEngine(
+                (
+                    f"Configuration error, cannot find engine {engine_name} "
+                    f"with version {engine_version}"
+                ),
+            )
+        return found_engine
+
+    def find_catalog(self, catalog_name: str) -> CatalogInfo:
+        """
+        Find a catalog defined in the server configuration
+        """
+        found_catalog = None
+        for catalog in self.catalogs:
+            if catalog.name == catalog_name:
+                found_catalog = catalog
+        if not found_catalog:
+            raise DJUnknownCatalog(
+                f"Configuration error, cannot find catalog {catalog_name}",
             )
-        catalog_entry = Catalog.parse_obj(catalog)
-        catalog_entry.engines = attached_engines
-        session.add(catalog_entry)
-    session.commit()
+        return found_catalog

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/db/postgres.py

@@ -0,0 +1,139 @@
+"""
+Dependency for getting the postgres pool and running backend DB queries
+"""
+
+# pylint: disable=too-many-arguments
+from datetime import datetime
+from typing import List
+from uuid import UUID
+
+from fastapi import Request
+from psycopg import sql
+from psycopg_pool import AsyncConnectionPool
+
+from djqs.exceptions import DJDatabaseError
+
+
+async def get_postgres_pool(request: Request) -> AsyncConnectionPool:
+    """
+    Get the postgres pool from the app instance
+    """
+    app = request.app
+    return app.state.pool
+
+
+class DBQuery:
+    """
+    Metadata DB queries using the psycopg composition utility
+    """
+
+    def __init__(self):
+        self._reset()
+
+    def _reset(self):
+        self.selects: List = []
+        self.inserts: List = []
+
+    def get_query(self, query_id: UUID):
+        """
+        Get metadata about a query
+        """
+        self.selects.append(
+            sql.SQL(
+                """
+                SELECT id, catalog_name, engine_name, engine_version, submitted_query,
+                       async_, executed_query, scheduled, started, finished, state, progress
+                FROM query
+                WHERE id = {query_id}
+                """,
+            ).format(query_id=sql.Literal(query_id)),
+        )
+        return self
+
+    def save_query(
+        self,
+        query_id: UUID,
+        catalog_name: str = "",
+        engine_name: str = "",
+        engine_version: str = "",
+        submitted_query: str = "",
+        async_: bool = False,
+        state: str = "",
+        progress: float = 0.0,
+        executed_query: str = None,
+        scheduled: datetime = None,
+        started: datetime = None,
+        finished: datetime = None,
+    ):
+        """
+        Save metadata about a query
+        """
+        self.inserts.append(
+            sql.SQL(
+                """
+                INSERT INTO query (id, catalog_name, engine_name, engine_version,
+                                   submitted_query, async_, executed_query, scheduled,
+                                   started, finished, state, progress)
+                VALUES ({query_id}, {catalog_name}, {engine_name}, {engine_version},
+                        {submitted_query}, {async_}, {executed_query}, {scheduled},
+                        {started}, {finished}, {state}, {progress})
+                ON CONFLICT (id) DO UPDATE SET
+                    catalog_name = EXCLUDED.catalog_name,
+                    engine_name = EXCLUDED.engine_name,
+                    engine_version = EXCLUDED.engine_version,
+                    submitted_query = EXCLUDED.submitted_query,
+                    async_ = EXCLUDED.async_,
+                    executed_query = EXCLUDED.executed_query,
+                    scheduled = EXCLUDED.scheduled,
+                    started = EXCLUDED.started,
+                    finished = EXCLUDED.finished,
+                    state = EXCLUDED.state,
+                    progress = EXCLUDED.progress
+                RETURNING *
+                """,
+            ).format(
+                query_id=sql.Literal(query_id),
+                catalog_name=sql.Literal(catalog_name),
+                engine_name=sql.Literal(engine_name),
+                engine_version=sql.Literal(engine_version),
+                submitted_query=sql.Literal(submitted_query),
+                async_=sql.Literal(async_),
+                executed_query=sql.Literal(executed_query),
+                scheduled=sql.Literal(scheduled),
+                started=sql.Literal(started),
+                finished=sql.Literal(finished),
+                state=sql.Literal(state),
+                progress=sql.Literal(progress),
+            ),
+        )
+        return self
+
+    async def execute(self, conn):
+        """
+        Submit all statements to the backend DB, multiple statements are submitted together
+        """
+        if not self.selects and not self.inserts:  # pragma: no cover
+            return
+
+        async with conn.cursor() as cur:
+            results = []
+            if len(self.inserts) > 1:  # pragma: no cover
+                async with conn.transaction():
+                    for statement in self.inserts:
+                        await cur.execute(statement)
+                results.append(await cur.fetchall())
+
+            if len(self.inserts) == 1:
+                await cur.execute(self.inserts[0])
+                if cur.rowcount == 0:  # pragma: no cover
+                    raise DJDatabaseError(
+                        "Insert statement resulted in no records being inserted",
+                    )
+                results.append((await cur.fetchone()))
+            if self.selects:
+                for statement in self.selects:
+                    await cur.execute(statement)
+                    results.append(await cur.fetchall())
+            await conn.commit()
+            self._reset()
+            return results

djqs/engine.py

@@ -2,28 +2,31 @@
 Query related functions.
 """
 
+import json
 import logging
 import os
-from datetime import datetime, timezone
-from typing import List, Tuple
+from dataclasses import asdict
+from datetime import date, datetime, timezone
+from typing import Dict, List, Optional, Tuple
 
 import duckdb
 import snowflake.connector
-import sqlparse
+from psycopg_pool import AsyncConnectionPool
 from sqlalchemy import create_engine, text
-from sqlmodel import Session, select
 
-from djqs.config import Settings
-from djqs.models.engine import Engine, EngineType
+from djqs.config import EngineType, Settings
+from djqs.constants import SQLALCHEMY_URI
+from djqs.db.postgres import DBQuery
+from djqs.exceptions import DJDatabaseError
 from djqs.models.query import (
     ColumnMetadata,
     Query,
     QueryResults,
     QueryState,
-    Results,
     StatementResults,
 )
 from djqs.typing import ColumnType, Description, SQLADialect, Stream, TypeEnum
+from djqs.utils import get_settings
 
 _logger = logging.getLogger(__name__)
 
@@ -66,9 +69,9 @@ def get_columns_from_description(
     return columns
 
 
-def run_query(
-    session: Session,
+def run_query(  # pylint: disable=R0914
     query: Query,
+    headers: Optional[Dict[str, str]] = None,
 ) -> List[Tuple[str, List[ColumnMetadata], Stream]]:
     """
     Run a query and return its results.
@@ -76,13 +79,26 @@ 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:
+
+    settings = get_settings()
+    engine_name = query.engine_name or settings.default_engine
+    engine_version = query.engine_version
+    engine = settings.find_engine(
+        engine_name=engine_name,
+        engine_version=engine_version,
+    )
+    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 +108,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,22 +117,21 @@ 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()
 
     output: List[Tuple[str, List[ColumnMetadata], Stream]] = []
-    statements = sqlparse.parse(query.executed_query)
-    for statement in statements:
-        # Druid doesn't like statements that end in a semicolon...
-        sql = str(statement).strip().rstrip(";")
-
-        results = connection.execute(text(sql))
-        stream = (tuple(row) for row in results)
-        columns = get_columns_from_description(
-            results.cursor.description,
-            sqla_engine.dialect,
-        )
-        output.append((sql, columns, stream))
+    results = connection.execute(text(query.executed_query))
+    stream = (tuple(row) for row in results)
+    columns = get_columns_from_description(
+        results.cursor.description,
+        sqla_engine.dialect,
+    )
+    output.append((query.executed_query, columns, stream))  # type: ignore
 
     return output
 
@@ -148,10 +164,24 @@ def run_snowflake_query(
     return output
 
 
-def process_query(
-    session: Session,
+def serialize_for_json(obj):
+    """
+    Handle serialization of date/datetimes for JSON output.
+    """
+    if isinstance(obj, list):
+        return [serialize_for_json(x) for x in obj]
+    if isinstance(obj, dict):
+        return {k: serialize_for_json(v) for k, v in obj.items()}
+    if isinstance(obj, (date, datetime)):
+        return obj.isoformat()
+    return obj
+
+
+async def process_query(
     settings: Settings,
+    postgres_pool: AsyncConnectionPool,
     query: Query,
+    headers: Optional[Dict[str, str]] = None,
 ) -> QueryResults:
     """
     Process a query.
@@ -163,10 +193,13 @@ def process_query(
     errors = []
     query.started = datetime.now(timezone.utc)
     try:
-        root = []
-        for sql, columns, stream in run_query(session=session, query=query):
+        results = []
+        for sql, columns, stream in run_query(
+            query=query,
+            headers=headers,
+        ):
             rows = list(stream)
-            root.append(
+            results.append(
                 StatementResults(
                     sql=sql,
                     columns=columns,
@@ -174,21 +207,52 @@ def process_query(
                     row_count=len(rows),
                 ),
             )
-        results = Results(__root__=root)
 
         query.state = QueryState.FINISHED
         query.progress = 1.0
     except Exception as ex:  # pylint: disable=broad-except
-        results = Results(__root__=[])
+        results = []
         query.state = QueryState.FAILED
         errors = [str(ex)]
 
     query.finished = datetime.now(timezone.utc)
 
-    session.add(query)
-    session.commit()
-    session.refresh(query)
+    async with postgres_pool.connection() as conn:
+        dbquery_results = (
+            await DBQuery()
+            .save_query(
+                query_id=query.id,
+                submitted_query=query.submitted_query,
+                state=QueryState.FINISHED.value,
+                async_=query.async_,
+            )
+            .execute(conn=conn)
+        )
+        query_save_result = dbquery_results[0]
+        if not query_save_result:  # pragma: no cover
+            raise DJDatabaseError("Query failed to save")
 
-    settings.results_backend.add(str(query.id), results.json())
+    settings.results_backend.add(
+        str(query.id),
+        json.dumps(
+            serialize_for_json(
+                [asdict(statement_result) for statement_result in results],
+            ),
+        ),
+    )
 
-    return QueryResults(results=results, errors=errors, **query.dict())
+    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,
+        scheduled=query.scheduled,
+        started=query.started,
+        finished=query.finished,
+        state=query.state,
+        progress=query.progress,
+        results=results,
+        errors=errors,
+    )

djqs/enum.py

@@ -1,6 +1,7 @@
 """
 Backwards-compatible StrEnum for both Python >= and < 3.11
 """
+
 import enum
 import sys