From b97d969f6d2fd79f3b2ef808f12c0ba201cb9d53 Mon Sep 17 00:00:00 2001 From: sadath-12 Date: Wed, 14 Aug 2024 14:38:11 +0530 Subject: [PATCH 1/2] rename tables in djqs Signed-off-by: sadath-12 --- datajunction-query/alembic/env.py | 6 +++--- datajunction-query/djqs/api/catalogs.py | 12 ++++++------ datajunction-query/djqs/api/engines.py | 6 +++--- datajunction-query/djqs/api/helpers.py | 10 +++++----- datajunction-query/djqs/api/queries.py | 6 +++--- datajunction-query/djqs/config.py | 18 +++++++++--------- datajunction-query/djqs/engine.py | 21 ++++++++++++--------- datajunction-query/djqs/models/catalog.py | 16 ++++++++-------- datajunction-query/djqs/models/engine.py | 2 +- datajunction-query/djqs/models/query.py | 2 +- 10 files changed, 51 insertions(+), 48 deletions(-) diff --git a/datajunction-query/alembic/env.py b/datajunction-query/alembic/env.py index ad293892e..d6421a7e2 100644 --- a/datajunction-query/alembic/env.py +++ b/datajunction-query/alembic/env.py @@ -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() diff --git a/datajunction-query/djqs/api/catalogs.py b/datajunction-query/djqs/api/catalogs.py index ded25f215..c10ed2d51 100644 --- a/datajunction-query/djqs/api/catalogs.py +++ b/datajunction-query/djqs/api/catalogs.py @@ -12,13 +12,13 @@ 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]) @@ -26,7 +26,7 @@ def list_catalogs(*, session: Session = Depends(get_session)) -> List[CatalogInf """ List all available catalogs """ - return list(session.exec(select(Catalog))) + return list(session.exec(select(QSCatalog))) @get_router.get("/catalogs/{name}/", response_model=CatalogInfo) @@ -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, @@ -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]: """ diff --git a/datajunction-query/djqs/api/engines.py b/datajunction-query/djqs/api/engines.py index 752a4ea60..aa9903f4a 100644 --- a/datajunction-query/djqs/api/engines.py +++ b/datajunction-query/djqs/api/engines.py @@ -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"]) @@ -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) @@ -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) diff --git a/datajunction-query/djqs/api/helpers.py b/datajunction-query/djqs/api/helpers.py index 3a1089a87..ffa9bc053 100644 --- a/datajunction-query/djqs/api/helpers.py +++ b/datajunction-query/djqs/api/helpers.py @@ -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( @@ -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 """ diff --git a/datajunction-query/djqs/api/queries.py b/datajunction-query/djqs/api/queries.py index 25253d9d1..7f0127cf8 100644 --- a/datajunction-query/djqs/api/queries.py +++ b/datajunction-query/djqs/api/queries.py @@ -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, @@ -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) @@ -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") diff --git a/datajunction-query/djqs/config.py b/datajunction-query/djqs/config.py index e68685d06..bdd7cd2c7 100644 --- a/datajunction-query/djqs/config.py +++ b/datajunction-query/djqs/config.py @@ -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 @@ -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" @@ -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: @@ -80,7 +80,7 @@ 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"]: @@ -88,9 +88,9 @@ def unknown_filetype(): 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() diff --git a/datajunction-query/djqs/engine.py b/datajunction-query/djqs/engine.py index 4bec0fdda..3c1924d01 100644 --- a/datajunction-query/djqs/engine.py +++ b/datajunction-query/djqs/engine.py @@ -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, @@ -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. @@ -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 @@ -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]]: """ @@ -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]]: """ @@ -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: """ diff --git a/datajunction-query/djqs/models/catalog.py b/datajunction-query/djqs/models/catalog.py index 10e231ea9..4d55408d4 100644 --- a/datajunction-query/djqs/models/catalog.py +++ b/datajunction-query/djqs/models/catalog.py @@ -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. """ @@ -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( diff --git a/datajunction-query/djqs/models/engine.py b/datajunction-query/djqs/models/engine.py index a790484fb..8eddd3e54 100644 --- a/datajunction-query/djqs/models/engine.py +++ b/datajunction-query/djqs/models/engine.py @@ -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. """ diff --git a/datajunction-query/djqs/models/query.py b/datajunction-query/djqs/models/query.py index 02c1ea2ba..7a3130514 100644 --- a/datajunction-query/djqs/models/query.py +++ b/datajunction-query/djqs/models/query.py @@ -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. """ From 89980f6854968752b35b53695afc1d384dc729a9 Mon Sep 17 00:00:00 2001 From: sadath-12 Date: Wed, 14 Aug 2024 14:58:46 +0530 Subject: [PATCH 2/2] correct docker command Signed-off-by: sadath-12 --- docker-compose.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker-compose.yml b/docker-compose.yml index 10858e621..97ca2e78c 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -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