Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions datajunction-query/alembic/env.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,9 @@
from sqlmodel import SQLModel, create_engine

from alembic import context
from djqs.models.catalog import Catalog
from djqs.models.engine import Engine
from djqs.models.query import Query
from djqs.models.catalog import QSCatalog
from djqs.models.engine import QSEngine
from djqs.models.query import QSQuery
from djqs.utils import get_settings

settings = get_settings()
Expand Down
12 changes: 6 additions & 6 deletions datajunction-query/djqs/api/catalogs.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,21 +12,21 @@
from djqs.api.engines import EngineInfo
from djqs.api.helpers import get_catalog, get_engine
from djqs.exceptions import DJException
from djqs.models.catalog import Catalog, CatalogInfo
from djqs.models.catalog import QSCatalog, CatalogInfo
from djqs.models.engine import BaseEngineInfo
from djqs.utils import get_session

_logger = logging.getLogger(__name__)
get_router = APIRouter(tags=["Catalogs & Engines"])
post_router = APIRouter(tags=["Catalogs & Engines - Dynamic Configuration"])
get_router = APIRouter(tags=["QSCatalogs & QSEngines"])
post_router = APIRouter(tags=["QSCatalogs & QSEngines - Dynamic Configuration"])


@get_router.get("/catalogs/", response_model=List[CatalogInfo])
def list_catalogs(*, session: Session = Depends(get_session)) -> List[CatalogInfo]:
"""
List all available catalogs
"""
return list(session.exec(select(Catalog)))
return list(session.exec(select(QSCatalog)))


@get_router.get("/catalogs/{name}/", response_model=CatalogInfo)
Expand Down Expand Up @@ -56,7 +56,7 @@ def add_catalog(
detail=f"Catalog already exists: `{data.name}`",
)

catalog = Catalog.from_orm(data)
catalog = QSCatalog.from_orm(data)
catalog.engines.extend(
list_new_engines(
session=session,
Expand Down Expand Up @@ -97,7 +97,7 @@ def add_engines_to_catalog(

def list_new_engines(
session: Session,
catalog: Catalog,
catalog: QSCatalog,
create_engines: List[EngineInfo],
) -> List[EngineInfo]:
"""
Expand Down
6 changes: 3 additions & 3 deletions datajunction-query/djqs/api/engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from sqlmodel import Session, select

from djqs.api.helpers import get_engine
from djqs.models.engine import BaseEngineInfo, Engine, EngineInfo
from djqs.models.engine import BaseEngineInfo, QSEngine, EngineInfo
from djqs.utils import get_session

get_router = APIRouter(tags=["Catalogs & Engines"])
Expand All @@ -21,7 +21,7 @@ def list_engines(*, session: Session = Depends(get_session)) -> List[EngineInfo]
"""
List all available engines
"""
return list(session.exec(select(Engine)))
return list(session.exec(select(QSEngine)))


@get_router.get("/engines/{name}/{version}/", response_model=BaseEngineInfo)
Expand Down Expand Up @@ -53,7 +53,7 @@ def add_engine(
detail=f"Engine already exists: `{data.name}` version `{data.version}`",
)

engine = Engine.from_orm(data)
engine = QSEngine.from_orm(data)
session.add(engine)
session.commit()
session.refresh(engine)
Expand Down
10 changes: 5 additions & 5 deletions datajunction-query/djqs/api/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,15 @@
from sqlmodel import Session, create_engine, select

from djqs.exceptions import DJException, DJTableNotFound
from djqs.models.catalog import Catalog
from djqs.models.engine import Engine
from djqs.models.catalog import QSCatalog
from djqs.models.engine import QSEngine


def get_catalog(session: Session, name: str) -> Catalog:
def get_catalog(session: Session, name: str) -> QSCatalog:
"""
Get a catalog by name
"""
statement = select(Catalog).where(Catalog.name == name)
statement = select(QSCatalog).where(QSCatalog.name == name)
catalog = session.exec(statement).one_or_none()
if not catalog:
raise DJException(
Expand All @@ -28,7 +28,7 @@ def get_catalog(session: Session, name: str) -> Catalog:
return catalog


def get_engine(session: Session, name: str, version: str) -> Engine:
def get_engine(session: Session, name: str, version: str) -> QSEngine:
"""
Return an Engine instance given an engine name and version
"""
Expand Down
6 changes: 3 additions & 3 deletions datajunction-query/djqs/api/queries.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
from djqs.config import Settings
from djqs.engine import process_query
from djqs.models.query import (
Query,
QSQuery,
QueryCreate,
QueryResults,
QueryState,
Expand Down Expand Up @@ -141,7 +141,7 @@ def save_query_and_run( # pylint: disable=R0913
"""
Store a new query to the DB and run it.
"""
query = Query(**create_query.dict(by_alias=True))
query = QSQuery(**create_query.dict(by_alias=True))
query.state = QueryState.ACCEPTED

session.add(query)
Expand Down Expand Up @@ -191,7 +191,7 @@ 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)
query = session.get(QSQuery, query_id)
if not query:
raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Query not found")

Expand Down
18 changes: 9 additions & 9 deletions datajunction-query/djqs/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
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.models.catalog import QSCatalog, QSCatalogEngines
from djqs.models.engine import QSEngine


class Settings(BaseSettings): # pylint: disable=too-few-public-methods
Expand All @@ -28,7 +28,7 @@ class Settings(BaseSettings): # pylint: disable=too-few-public-methods
url: str = "http://localhost:8001/"

# SQLAlchemy URI for the metadata database.
index: str = "sqlite:///djqs.db?check_same_thread=False"
index: str = "postgresql+psycopg://dj:dj@postgres_metadata:5432/dj"

# The default engine to use for reflection
default_reflection_engine: str = "default"
Expand Down Expand Up @@ -59,9 +59,9 @@ def load_djqs_config(settings: Settings, session: Session) -> None: # pragma: n
if not config_file:
return

session.exec(delete(Catalog))
session.exec(delete(Engine))
session.exec(delete(CatalogEngines))
session.exec(delete(QSCatalog))
session.exec(delete(QSEngine))
session.exec(delete(QSCatalogEngines))
session.commit()

with open(config_file, mode="r", encoding="utf-8") as filestream:
Expand All @@ -80,17 +80,17 @@ def unknown_filetype():
)

for engine in data["engines"]:
session.add(Engine.parse_obj(engine))
session.add(QSEngine.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(),
session.exec(select(QSEngine).where(QSEngine.name == name)).one(),
)
catalog_entry = Catalog.parse_obj(catalog)
catalog_entry = QSCatalog.parse_obj(catalog)
catalog_entry.engines = attached_engines
session.add(catalog_entry)
session.commit()
21 changes: 12 additions & 9 deletions datajunction-query/djqs/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,13 @@
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.engine import QSEngine, EngineType

from djqs.models.query import (
ColumnMetadata,
Query,
QSQuery,
QueryResults,
QueryState,
Results,
Expand Down Expand Up @@ -68,8 +70,9 @@ def get_columns_from_description(

def run_query( # pylint: disable=R0914
session: Session,
query: Query,
query: QSQuery,
headers: Optional[Dict[str, str]] = None,

) -> List[Tuple[str, List[ColumnMetadata], Stream]]:
"""
Run a query and return its results.
Expand All @@ -81,9 +84,9 @@ def run_query( # pylint: disable=R0914
_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),
select(QSEngine)
.where(QSEngine.name == query.engine_name)
.where(QSEngine.version == query.engine_version),
).one()

query_server = headers.get("SQLALCHEMY_URI") if headers else None
Expand Down Expand Up @@ -139,7 +142,7 @@ def run_query( # pylint: disable=R0914


def run_duckdb_query(
query: Query,
query: QSQuery,
conn: duckdb.DuckDBPyConnection,
) -> List[Tuple[str, List[ColumnMetadata], Stream]]:
"""
Expand All @@ -153,7 +156,7 @@ def run_duckdb_query(


def run_snowflake_query(
query: Query,
query: QSQuery,
cur: snowflake.connector.cursor.SnowflakeCursor,
) -> List[Tuple[str, List[ColumnMetadata], Stream]]:
"""
Expand All @@ -169,7 +172,7 @@ def run_snowflake_query(
def process_query(
session: Session,
settings: Settings,
query: Query,
query: QSQuery,
headers: Optional[Dict[str, str]] = None,
) -> QueryResults:
"""
Expand Down
16 changes: 8 additions & 8 deletions datajunction-query/djqs/models/catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,13 @@
from sqlalchemy_utils import UUIDType
from sqlmodel import JSON, Field, Relationship, SQLModel

from djqs.models.engine import BaseEngineInfo, Engine
from djqs.models.engine import BaseEngineInfo, QSEngine

if TYPE_CHECKING:
from djqs.utils import UTCDatetime


class CatalogEngines(SQLModel, table=True): # type: ignore
class QSCatalogEngines(SQLModel, table=True): # type: ignore
"""
Join table for catalogs and engines.
"""
Expand All @@ -34,19 +34,19 @@ class CatalogEngines(SQLModel, table=True): # type: ignore
)


class Catalog(SQLModel, table=True): # type: ignore
class QSCatalog(SQLModel, table=True): # type: ignore
"""
A catalog.
A QScatalog.
"""

id: Optional[int] = Field(default=None, primary_key=True)
uuid: UUID = Field(default_factory=uuid4, sa_column=SqlaColumn(UUIDType()))
name: str
engines: List[Engine] = Relationship(
link_model=CatalogEngines,
engines: List[QSEngine] = Relationship(
link_model=QSCatalogEngines,
sa_relationship_kwargs={
"primaryjoin": "Catalog.id==CatalogEngines.catalog_id",
"secondaryjoin": "Engine.id==CatalogEngines.engine_id",
"primaryjoin": "Catalog.id==QSCatalogEngines.catalog_id",
"secondaryjoin": "Engine.id==QSCatalogEngines.engine_id",
},
)
created_at: "UTCDatetime" = Field(
Expand Down
2 changes: 1 addition & 1 deletion datajunction-query/djqs/models/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ class EngineType(Enum):
SNOWFLAKE = "snowflake"


class Engine(SQLModel, table=True): # type: ignore
class QSEngine(SQLModel, table=True): # type: ignore
"""
A query engine.
"""
Expand Down
2 changes: 1 addition & 1 deletion datajunction-query/djqs/models/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ class Config: # pylint: disable=too-few-public-methods, missing-class-docstring
allow_population_by_field_name = True


class Query(BaseQuery, table=True): # type: ignore
class QSQuery(BaseQuery, table=True): # type: ignore
"""
A query.
"""
Expand Down
2 changes: 1 addition & 1 deletion docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ services:
context: ./datajunction-query
volumes:
- ./datajunction-query:/code
command: /bin/bash -c "cp -n /code/djqs.demo.db /code/djqs.db | true && sleep 2 && alembic upgrade head"
command: /bin/bash -c "alembic upgrade head"
restart: on-failure
depends_on:
- djqs
Expand Down