datajunction-query 0.0.1a1.dev0__tar.gz → 0.0.1a31__tar.gz

{datajunction_query-0.0.1a1.dev0 → datajunction_query-0.0.1a31}/.pre-commit-config.yaml

@@ -72,6 +72,7 @@ repos:
     - types-pkg_resources
     - types-PyYAML
     - types-tabulate
+    - types-toml
 - repo: https://github.com/asottile/add-trailing-comma
   rev: v2.2.1
   hooks:

{datajunction_query-0.0.1a1.dev0 → datajunction_query-0.0.1a31}/Makefile

@@ -47,4 +47,4 @@ lint:
 dev-release:
 	hatch version dev
 	hatch build
-	hatch publish
+	hatch publish

{datajunction_query-0.0.1a1.dev0 → datajunction_query-0.0.1a31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datajunction-query
-Version: 0.0.1a1.dev0
+Version: 0.0.1a31
 Summary: OSS Implementation of a DataJunction Query Service
 Project-URL: repository, https://github.com/DataJunction/dj
 Author-email: DataJunction Authors <roberto@dealmeida.net>
@@ -30,6 +30,7 @@ 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: toml>=0.10.2
 Provides-Extra: uvicorn
 Requires-Dist: uvicorn[standard]>=0.21.1; extra == 'uvicorn'
 Description-Content-Type: text/x-rst

datajunction_query-0.0.1a31/alembic/versions/2023_10_09_1858-f3407a1ec625_add_type_and_extra_para_ms_field_for_.py

@@ -0,0 +1,36 @@
+"""Add type and extra_params field for engines
+
+Revision ID: f3407a1ec625
+Revises: a7e11a2438b4
+Create Date: 2023-10-09 18:58:57.462134+00:00
+
+"""
+# pylint: disable=no-member, invalid-name, missing-function-docstring, unused-import, no-name-in-module
+
+import sqlalchemy as sa
+import sqlmodel
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "f3407a1ec625"
+down_revision = "a7e11a2438b4"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("engine", sa.Column("extra_params", sa.JSON(), nullable=True))
+    op.add_column(
+        "engine",
+        sa.Column("type", sa.Enum("DUCKDB", name="enginetype"), nullable=False),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("engine", "type")
+    op.drop_column("engine", "extra_params")
+    # ### end Alembic commands ###

datajunction_query-0.0.1a31/config.djqs.yml

@@ -0,0 +1,11 @@
+engines:
+  - name: duckdb
+    version: 0.7.1
+    type: duckdb
+    uri: duckdb:////code/docker/default.duckdb
+    extra_params:
+      location: /code/docker/default.duckdb
+catalogs:
+  - name: warehouse
+    engines:
+      - duckdb

datajunction_query-0.0.1a31/config.jsonschema

@@ -0,0 +1,83 @@
+{
+    "$schema": "http://json-schema.org/draft-06/schema#",
+    "$ref": "#/definitions/DJQSConfig",
+    "definitions": {
+        "DJQSConfig": {
+            "type": "object",
+            "additionalProperties": false,
+            "properties": {
+                "engines": {
+                    "type": "array",
+                    "items": {
+                        "$ref": "#/definitions/Engine"
+                    }
+                },
+                "catalogs": {
+                    "type": "array",
+                    "items": {
+                        "$ref": "#/definitions/Catalog"
+                    }
+                }
+            },
+            "required": [
+                "catalogs",
+                "engines"
+            ],
+            "title": "DJQSConfig"
+        },
+        "Catalog": {
+            "type": "object",
+            "additionalProperties": false,
+            "properties": {
+                "name": {
+                    "type": "string"
+                },
+                "engines": {
+                    "type": "array",
+                    "items": {
+                        "type": "string"
+                    }
+                }
+            },
+            "required": [
+                "engines",
+                "name"
+            ],
+            "title": "Catalog"
+        },
+        "Engine": {
+            "type": "object",
+            "additionalProperties": false,
+            "properties": {
+                "name": {
+                    "type": "string"
+                },
+                "version": {
+                    "type": "string"
+                },
+                "type": {
+                    "type": "string"
+                },
+                "uri": {
+                    "type": "string"
+                },
+                "extra_params": {
+                    "$ref": "#/definitions/ExtraParams"
+                }
+            },
+            "required": [
+                "extra_params",
+                "name",
+                "type",
+                "uri",
+                "version"
+            ],
+            "title": "Engine"
+        },
+        "ExtraParams": {
+            "type": "object",
+            "additionalProperties": true,
+            "title": "ExtraParams"
+        }
+    }
+}

datajunction_query-0.0.1a31/djqs/__about__.py

@@ -0,0 +1,4 @@
+"""
+Version for Hatch
+"""
+__version__ = "0.0.1a31"

{datajunction_query-0.0.1a1.dev0 → datajunction_query-0.0.1a31}/djqs/api/catalogs.py

@@ -17,10 +17,11 @@ from djqs.models.engine import BaseEngineInfo
 from djqs.utils import get_session
 
 _logger = logging.getLogger(__name__)
-router = APIRouter(tags=["Catalogs & Engines"])
+get_router = APIRouter(tags=["Catalogs & Engines"])
+post_router = APIRouter(tags=["Catalogs & Engines - Dynamic Configuration"])
 
 
-@router.get("/catalogs/", response_model=List[CatalogInfo])
+@get_router.get("/catalogs/", response_model=List[CatalogInfo])
 def list_catalogs(*, session: Session = Depends(get_session)) -> List[CatalogInfo]:
     """
     List all available catalogs
@@ -28,7 +29,7 @@ def list_catalogs(*, session: Session = Depends(get_session)) -> List[CatalogInf
     return list(session.exec(select(Catalog)))
 
 
-@router.get("/catalogs/{name}/", response_model=CatalogInfo)
+@get_router.get("/catalogs/{name}/", response_model=CatalogInfo)
 def read_catalog(name: str, *, session: Session = Depends(get_session)) -> CatalogInfo:
     """
     Return a catalog by name
@@ -36,7 +37,7 @@ def read_catalog(name: str, *, session: Session = Depends(get_session)) -> Catal
     return get_catalog(session, name)
 
 
-@router.post("/catalogs/", response_model=CatalogInfo, status_code=201)
+@post_router.post("/catalogs/", response_model=CatalogInfo, status_code=201)
 def add_catalog(
     data: CatalogInfo,
     *,
@@ -70,7 +71,11 @@ def add_catalog(
     return catalog
 
 
-@router.post("/catalogs/{name}/engines/", response_model=CatalogInfo, status_code=201)
+@post_router.post(
+    "/catalogs/{name}/engines/",
+    response_model=CatalogInfo,
+    status_code=201,
+)
 def add_engines_to_catalog(
     name: str,
     data: List[BaseEngineInfo],

{datajunction_query-0.0.1a1.dev0 → datajunction_query-0.0.1a31}/djqs/api/engines.py

@@ -12,10 +12,11 @@ from djqs.api.helpers import get_engine
 from djqs.models.engine import BaseEngineInfo, Engine, EngineInfo
 from djqs.utils import get_session
 
-router = APIRouter(tags=["Catalogs & Engines"])
+get_router = APIRouter(tags=["Catalogs & Engines"])
+post_router = APIRouter(tags=["Catalogs & Engines - Dynamic Configuration"])
 
 
-@router.get("/engines/", response_model=List[EngineInfo])
+@get_router.get("/engines/", response_model=List[EngineInfo])
 def list_engines(*, session: Session = Depends(get_session)) -> List[EngineInfo]:
     """
     List all available engines
@@ -23,7 +24,7 @@ def list_engines(*, session: Session = Depends(get_session)) -> List[EngineInfo]
     return list(session.exec(select(Engine)))
 
 
-@router.get("/engines/{name}/{version}/", response_model=BaseEngineInfo)
+@get_router.get("/engines/{name}/{version}/", response_model=BaseEngineInfo)
 def list_engine(
     name: str, version: str, *, session: Session = Depends(get_session)
 ) -> BaseEngineInfo:
@@ -33,7 +34,7 @@ def list_engine(
     return get_engine(session, name, version)
 
 
-@router.post("/engines/", response_model=BaseEngineInfo, status_code=201)
+@post_router.post("/engines/", response_model=BaseEngineInfo, status_code=201)
 def add_engine(
     data: EngineInfo,
     *,

{datajunction_query-0.0.1a1.dev0 → datajunction_query-0.0.1a31}/djqs/api/helpers.py

@@ -9,7 +9,6 @@ from sqlalchemy import inspect
 from sqlalchemy.exc import NoResultFound, NoSuchTableError, OperationalError
 from sqlmodel import Session, create_engine, select
 
-from djqs.engine import describe_table_via_spark, get_spark_session
 from djqs.exceptions import DJException
 from djqs.models.catalog import Catalog
 from djqs.models.engine import Engine
@@ -59,14 +58,6 @@ def get_columns(
     if not uri:
         raise DJException("Cannot retrieve columns without a uri")
 
-    if uri.startswith("spark://"):
-        spark = get_spark_session()
-        return describe_table_via_spark(
-            spark,
-            schema,
-            table,
-        )
-
     engine = create_engine(uri, **extra_params)
     try:
         inspector = inspect(engine)

{datajunction_query-0.0.1a1.dev0 → datajunction_query-0.0.1a31}/djqs/api/main.py

@@ -5,7 +5,7 @@ Main DJ query server app.
 # All the models need to be imported here so that SQLModel can define their
 # relationships at runtime without causing circular imports.
 # See https://sqlmodel.tiangolo.com/tutorial/code-structure/#make-circular-imports-work.
-# pylint: disable=unused-import
+# pylint: disable=unused-import,expression-not-assigned
 
 import logging
 
@@ -14,13 +14,16 @@ from fastapi.responses import JSONResponse
 
 from djqs import __version__
 from djqs.api import catalogs, engines, queries, tables
+from djqs.config import load_djqs_config
 from djqs.exceptions import DJException
-from djqs.utils import get_settings
+from djqs.utils import get_session, get_settings
 
 _logger = logging.getLogger(__name__)
 
-
 settings = get_settings()
+session = next(get_session())
+load_djqs_config(settings=settings, session=session)
+
 app = FastAPI(
     title=settings.name,
     description=settings.description,
@@ -30,10 +33,12 @@ app = FastAPI(
         "url": "https://mit-license.org/",
     },
 )
-app.include_router(catalogs.router)
-app.include_router(engines.router)
+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.exception_handler(DJException)

datajunction_query-0.0.1a31/djqs/config.py

@@ -0,0 +1,96 @@
+"""
+Configuration for the query service
+"""
+
+import json
+from datetime import timedelta
+from typing import 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
+
+
+class Settings(BaseSettings):  # pylint: disable=too-few-public-methods
+    """
+    Configuration for the query service
+    """
+
+    name: str = "DJQS"
+    description: str = "A DataJunction Query Service"
+    url: str = "http://localhost:8001/"
+
+    # 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"
+
+    # The default engine version to use for reflection
+    default_reflection_engine_version: str = ""
+
+    # Where to store the results from queries.
+    results_backend: BaseCache = FileSystemCache("/tmp/djqs", default_timeout=0)
+
+    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)
+
+    # Configuration file for catalogs and engines
+    configuration_file: Optional[str] = None
+
+    # Enable setting catalog and engine config via REST API calls
+    enable_dynamic_config: bool = True
+
+
+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()
+        )
+
+    for engine in data["engines"]:
+        session.add(Engine.parse_obj(engine))
+    session.commit()
+
+    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(),
+            )
+        catalog_entry = Catalog.parse_obj(catalog)
+        catalog_entry.engines = attached_engines
+        session.add(catalog_entry)
+    session.commit()

{datajunction_query-0.0.1a1.dev0 → datajunction_query-0.0.1a31}/djqs/engine.py

@@ -4,17 +4,16 @@ Query related functions.
 
 import logging
 from datetime import datetime, timezone
-from typing import List, Optional, Tuple
+from typing import List, Tuple
 
 import duckdb
 import sqlparse
-from pyspark.sql import SparkSession  # pylint: disable=import-error
 from sqlalchemy import create_engine, text
 from sqlmodel import Session, select
 
 from djqs.config import Settings
 from djqs.models.catalog import Catalog
-from djqs.models.engine import Engine
+from djqs.models.engine import Engine, EngineType
 from djqs.models.query import (
     ColumnMetadata,
     Query,
@@ -85,11 +84,15 @@ def run_query(
         .where(Engine.name == query.engine_name)
         .where(Engine.version == query.engine_version),
     ).one()
-    if engine.uri == "spark://local[*]":
-        spark = get_spark_session()
-        return run_spark_query(query, spark)
-    if engine.uri == "duckdb://local[*]":
-        conn = duckdb.connect(database="/code/docker/default.duckdb", read_only=False)
+    if engine.type == EngineType.DUCKDB:
+        conn = (
+            duckdb.connect()
+            if engine.uri == "duckdb:///:memory:"
+            else duckdb.connect(
+                database=engine.extra_params["location"],
+                read_only=True,
+            )
+        )
         return run_duckdb_query(query, conn)
     sqla_engine = create_engine(engine.uri, **catalog.extra_params)
     connection = sqla_engine.connect()
@@ -111,49 +114,6 @@ def run_query(
     return output
 
 
-def get_spark_session():
-    """
-    Get a spark session
-    """
-    SparkSession._instantiatedContext = None  # pylint: disable=protected-access
-    spark = (
-        SparkSession.builder.master("local[*]")
-        .appName("djqs")
-        .enableHiveSupport()
-        .getOrCreate()
-    )
-    return spark
-
-
-def run_spark_query(
-    query: Query,
-    spark: SparkSession,
-) -> List[Tuple[str, List[ColumnMetadata], Stream]]:
-    """
-    Run a spark SQL query against the local warehouse
-    """
-    output: List[Tuple[str, List[ColumnMetadata], Stream]] = []
-    results_df = spark.sql(query.submitted_query)
-    rows = results_df.rdd.map(tuple).collect()
-    columns: List[ColumnMetadata] = []
-    output.append((query.submitted_query, columns, rows))
-    return output
-
-
-def describe_table_via_spark(
-    spark: SparkSession,
-    schema: Optional[str],
-    table: str,
-):
-    """
-    Gets the column schemas.
-    """
-    schema_ = f"{schema}." if schema else ""
-    schema_df = spark.sql(f"DESCRIBE TABLE {schema_}{table};")
-    rows = schema_df.rdd.map(tuple).collect()
-    return [{"name": row[0], "type": row[1]} for row in rows]
-
-
 def run_duckdb_query(
     query: Query,
     conn: duckdb.DuckDBPyConnection,

{datajunction_query-0.0.1a1.dev0 → datajunction_query-0.0.1a31}/djqs/models/engine.py

@@ -1,10 +1,20 @@
 """
 Models for columns.
 """
+from enum import Enum
+from typing import Dict, Optional
 
-from typing import Optional
+from sqlalchemy.sql.schema import Column as SqlaColumn
+from sqlmodel import JSON, Field, SQLModel
 
-from sqlmodel import Field, SQLModel
+
+class EngineType(Enum):
+    """
+    Supported engine types
+    """
+
+    DUCKDB = "duckdb"
+    SQLALCHEMY = "sqlalchemy"
 
 
 class Engine(SQLModel, table=True):  # type: ignore
@@ -14,8 +24,10 @@ class Engine(SQLModel, table=True):  # type: ignore
 
     id: Optional[int] = Field(default=None, primary_key=True)
     name: str
+    type: EngineType
     version: str
     uri: Optional[str]
+    extra_params: Dict = Field(default={}, sa_column=SqlaColumn(JSON))
 
 
 class BaseEngineInfo(SQLModel):
@@ -25,6 +37,8 @@ class BaseEngineInfo(SQLModel):
 
     name: str
     version: str
+    type: EngineType
+    extra_params: Dict = {}
 
 
 class EngineInfo(BaseEngineInfo):