From 5fe82bc3c113bd06bd34fa0d2006762a4b2a0d25 Mon Sep 17 00:00:00 2001 From: Grieve Date: Fri, 21 Jun 2024 10:54:22 +0800 Subject: [PATCH] Integrate API of ibis (#624) * Integrate API * Add isort into linter list --- ibis-server/app/dependencies.py | 7 ++ ibis-server/app/model/__init__.py | 91 ++++++++++++++ ibis-server/app/model/connector.py | 43 +------ ibis-server/app/model/data_source.py | 111 +++++++----------- ibis-server/app/model/metadata/bigquery.py | 8 +- ibis-server/app/model/metadata/dto.py | 4 +- ibis-server/app/model/metadata/factory.py | 11 +- ibis-server/app/model/metadata/metadata.py | 4 +- ibis-server/app/model/metadata/mysql.py | 9 +- ibis-server/app/model/metadata/postgres.py | 9 +- ibis-server/app/model/validator.py | 9 -- ibis-server/app/routers/v2/__init__.py | 2 +- ibis-server/app/routers/v2/analysis.py | 7 +- ibis-server/app/routers/v2/ibis.py | 57 +++++++++ ibis-server/app/routers/v2/ibis/__init__.py | 12 -- ibis-server/app/routers/v2/ibis/bigquery.py | 49 -------- ibis-server/app/routers/v2/ibis/mysql.py | 50 -------- ibis-server/app/routers/v2/ibis/postgres.py | 49 -------- ibis-server/app/routers/v2/ibis/snowflake.py | 33 ------ ibis-server/pyproject.toml | 3 + ibis-server/tests/routers/ibis/test_mysql.py | 4 +- .../tests/routers/ibis/test_postgres.py | 2 +- 22 files changed, 233 insertions(+), 341 deletions(-) create mode 100644 ibis-server/app/dependencies.py create mode 100644 ibis-server/app/routers/v2/ibis.py delete mode 100644 ibis-server/app/routers/v2/ibis/__init__.py delete mode 100644 ibis-server/app/routers/v2/ibis/bigquery.py delete mode 100644 ibis-server/app/routers/v2/ibis/mysql.py delete mode 100644 ibis-server/app/routers/v2/ibis/postgres.py delete mode 100644 ibis-server/app/routers/v2/ibis/snowflake.py diff --git a/ibis-server/app/dependencies.py b/ibis-server/app/dependencies.py new file mode 100644 index 000000000..65c9d8adc --- /dev/null +++ b/ibis-server/app/dependencies.py @@ -0,0 +1,7 @@ +from app.model import QueryDTO +from app.model.data_source import DataSource + + +# Rebuild model to validate the dto is correct via validation of the pydantic +def verify_query_dto(data_source: DataSource, dto: QueryDTO): + data_source.get_dto_type()(**dto.model_dump(by_alias=True)) diff --git a/ibis-server/app/model/__init__.py b/ibis-server/app/model/__init__.py index e69de29bb..bd5ff7f59 100644 --- a/ibis-server/app/model/__init__.py +++ b/ibis-server/app/model/__init__.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from typing import Union + +from pydantic import BaseModel, Field + +manifest_str_field = Field(alias="manifestStr", description="Base64 manifest") +connection_info_field = Field(alias="connectionInfo") + + +class QueryDTO(BaseModel): + sql: str + manifest_str: str = manifest_str_field + column_dtypes: dict[str, str] | None = Field( + alias="columnDtypes", + description="If this field is set, it will forcibly convert the type.", + default=None, + ) + connection_info: ConnectionInfo = connection_info_field + + +class QueryBigQueryDTO(QueryDTO): + connection_info: BigQueryConnectionInfo = connection_info_field + + +class QueryMySqlDTO(QueryDTO): + connection_info: ConnectionUrl | MySqlConnectionInfo = connection_info_field + + +class QueryPostgresDTO(QueryDTO): + connection_info: ConnectionUrl | PostgresConnectionInfo = connection_info_field + + +class QuerySnowflakeDTO(QueryDTO): + connection_info: SnowflakeConnectionInfo = connection_info_field + + +class BigQueryConnectionInfo(BaseModel): + project_id: str + dataset_id: str + credentials: str = Field(description="Base64 encode `credentials.json`") + + +class MySqlConnectionInfo(BaseModel): + host: str + port: int + database: str + user: str + password: str + + +class ConnectionUrl(BaseModel): + connection_url: str = Field(alias="connectionUrl") + + +class PostgresConnectionInfo(BaseModel): + host: str = Field(examples=["localhost"]) + port: int = Field(examples=[5432]) + database: str + user: str + password: str + + +class SnowflakeConnectionInfo(BaseModel): + user: str + password: str + account: str + database: str + sf_schema: str = Field( + alias="schema" + ) # Use `sf_schema` to avoid `schema` shadowing in BaseModel + + +ConnectionInfo = Union[ + BigQueryConnectionInfo, + ConnectionUrl, + MySqlConnectionInfo, + PostgresConnectionInfo, + SnowflakeConnectionInfo, +] + + +class ValidateDTO(BaseModel): + manifest_str: str = manifest_str_field + parameters: dict[str, str] + connection_info: ConnectionInfo = connection_info_field + + +class AnalyzeSQLDTO(BaseModel): + manifest_str: str = manifest_str_field + sql: str diff --git a/ibis-server/app/model/connector.py b/ibis-server/app/model/connector.py index bd686acee..b60d258b5 100644 --- a/ibis-server/app/model/connector.py +++ b/ibis-server/app/model/connector.py @@ -1,19 +1,10 @@ from json import loads import pandas as pd -from pydantic import BaseModel, Field from app.mdl.rewriter import rewrite -from app.model.data_source import ( - BigQueryConnectionInfo, - ConnectionInfo, - DataSource, - MySqlConnectionUrl, - MySqlConnectionInfo, - PostgresConnectionUrl, - PostgresConnectionInfo, - SnowflakeConnectionInfo, -) +from app.model import ConnectionInfo +from app.model.data_source import DataSource class Connector: @@ -67,35 +58,5 @@ def _to_datetime_and_format(series: pd.Series) -> pd.Series: ) -class QueryDTO(BaseModel): - sql: str - manifest_str: str = Field(alias="manifestStr", description="Base64 manifest") - column_dtypes: dict[str, str] | None = Field( - alias="columnDtypes", - description="If this field is set, it will forcibly convert the type.", - default=None, - ) - - -class QueryBigQueryDTO(QueryDTO): - connection_info: BigQueryConnectionInfo = Field(alias="connectionInfo") - - -class QueryMySqlDTO(QueryDTO): - connection_info: MySqlConnectionUrl | MySqlConnectionInfo = Field( - alias="connectionInfo" - ) - - -class QueryPostgresDTO(QueryDTO): - connection_info: PostgresConnectionUrl | PostgresConnectionInfo = Field( - alias="connectionInfo" - ) - - -class QuerySnowflakeDTO(QueryDTO): - connection_info: SnowflakeConnectionInfo = Field(alias="connectionInfo") - - class QueryDryRunError(Exception): pass diff --git a/ibis-server/app/model/data_source.py b/ibis-server/app/model/data_source.py index 4af62bd11..f99570b88 100644 --- a/ibis-server/app/model/data_source.py +++ b/ibis-server/app/model/data_source.py @@ -1,14 +1,26 @@ from __future__ import annotations import base64 -from enum import StrEnum, auto +from enum import Enum, StrEnum, auto from json import loads -from typing import Union import ibis from google.oauth2 import service_account from ibis import BaseBackend -from pydantic import BaseModel, Field + +from app.model import ( + BigQueryConnectionInfo, + ConnectionInfo, + ConnectionUrl, + MySqlConnectionInfo, + PostgresConnectionInfo, + QueryBigQueryDTO, + QueryDTO, + QueryMySqlDTO, + QueryPostgresDTO, + QuerySnowflakeDTO, + SnowflakeConnectionInfo, +) class DataSource(StrEnum): @@ -17,18 +29,33 @@ class DataSource(StrEnum): postgres = auto() snowflake = auto() - def get_connection(self, dto) -> BaseBackend: - match self: - case DataSource.bigquery: - return self.get_bigquery_connection(dto) - case DataSource.mysql: - return self.get_mysql_connection(dto) - case DataSource.postgres: - return self.get_postgres_connection(dto) - case DataSource.snowflake: - return self.get_snowflake_connection(dto) - case _: - raise NotImplementedError(f"Unsupported data source: {self}") + def get_connection(self, info: ConnectionInfo) -> BaseBackend: + try: + return DataSourceExtension[self].get_connection(info) + except KeyError: + raise NotImplementedError(f"Unsupported data source: {self}") + + def get_dto_type(self): + try: + return DataSourceExtension[self].dto + except KeyError: + raise NotImplementedError(f"Unsupported data source: {self}") + + +class DataSourceExtension(Enum): + bigquery = QueryBigQueryDTO + mysql = QueryMySqlDTO + postgres = QueryPostgresDTO + snowflake = QuerySnowflakeDTO + + def __init__(self, dto: QueryDTO): + self.dto = dto + + def get_connection(self, info: ConnectionInfo) -> BaseBackend: + try: + return getattr(self, f"get_{self.name}_connection")(info) + except KeyError: + raise NotImplementedError(f"Unsupported data source: {self}") @staticmethod def get_bigquery_connection(info: BigQueryConnectionInfo) -> BaseBackend: @@ -44,7 +71,7 @@ def get_bigquery_connection(info: BigQueryConnectionInfo) -> BaseBackend: @staticmethod def get_mysql_connection( - info: MySqlConnectionUrl | MySqlConnectionInfo, + info: ConnectionUrl | MySqlConnectionInfo, ) -> BaseBackend: return ibis.connect( getattr(info, "connection_url", None) @@ -54,7 +81,7 @@ def get_mysql_connection( @staticmethod def get_postgres_connection( - info: PostgresConnectionUrl | PostgresConnectionInfo, + info: ConnectionUrl | PostgresConnectionInfo, ) -> BaseBackend: return ibis.connect( getattr(info, "connection_url", None) @@ -70,53 +97,3 @@ def get_snowflake_connection(info: SnowflakeConnectionInfo) -> BaseBackend: database=info.database, schema=info.sf_schema, ) - - -class BigQueryConnectionInfo(BaseModel): - project_id: str - dataset_id: str - credentials: str = Field(description="Base64 encode `credentials.json`") - - -class MySqlConnectionInfo(BaseModel): - host: str - port: int - database: str - user: str - password: str - - -class MySqlConnectionUrl(BaseModel): - connection_url: str = Field(alias="connectionUrl") - - -class PostgresConnectionInfo(BaseModel): - host: str = Field(examples=["localhost"]) - port: int = Field(examples=[5432]) - database: str - user: str - password: str - - -class PostgresConnectionUrl(BaseModel): - connection_url: str = Field(alias="connectionUrl") - - -class SnowflakeConnectionInfo(BaseModel): - user: str - password: str - account: str - database: str - sf_schema: str = Field( - alias="schema" - ) # Use `sf_schema` to avoid `schema` shadowing in BaseModel - - -ConnectionInfo = Union[ - BigQueryConnectionInfo, - MySqlConnectionInfo, - MySqlConnectionUrl, - PostgresConnectionInfo, - PostgresConnectionUrl, - SnowflakeConnectionInfo, -] diff --git a/ibis-server/app/model/metadata/bigquery.py b/ibis-server/app/model/metadata/bigquery.py index d5c78ef31..277cf3ca5 100644 --- a/ibis-server/app/model/metadata/bigquery.py +++ b/ibis-server/app/model/metadata/bigquery.py @@ -1,13 +1,13 @@ from json import loads -from app.model.data_source import BigQueryConnectionInfo +from app.model import BigQueryConnectionInfo from app.model.data_source import DataSource from app.model.metadata.dto import ( - Table, - Constraint, - TableProperties, Column, + Constraint, ConstraintType, + Table, + TableProperties, ) from app.model.metadata.metadata import Metadata diff --git a/ibis-server/app/model/metadata/dto.py b/ibis-server/app/model/metadata/dto.py index 7c543be44..e7865e1cf 100644 --- a/ibis-server/app/model/metadata/dto.py +++ b/ibis-server/app/model/metadata/dto.py @@ -1,9 +1,9 @@ from enum import Enum -from typing import List, Optional, Dict, Any +from typing import Any, Dict, List, Optional from pydantic import BaseModel, Field -from app.model.data_source import ConnectionInfo +from app.model import ConnectionInfo class MetadataDTO(BaseModel): diff --git a/ibis-server/app/model/metadata/factory.py b/ibis-server/app/model/metadata/factory.py index dae843217..66fffe813 100644 --- a/ibis-server/app/model/metadata/factory.py +++ b/ibis-server/app/model/metadata/factory.py @@ -1,14 +1,15 @@ from json import loads -from app.model.data_source import DataSource, ConnectionInfo +from app.model import ConnectionInfo +from app.model.data_source import DataSource from app.model.metadata.bigquery import BigQueryMetadata -from app.model.metadata.metadata import Metadata -from app.model.metadata.postgres import PostgresMetadata -from app.model.metadata.mysql import MySQLMetadata from app.model.metadata.dto import ( - Table, Constraint, + Table, ) +from app.model.metadata.metadata import Metadata +from app.model.metadata.mysql import MySQLMetadata +from app.model.metadata.postgres import PostgresMetadata class MetadataFactory: diff --git a/ibis-server/app/model/metadata/metadata.py b/ibis-server/app/model/metadata/metadata.py index cde79fac0..021607d58 100644 --- a/ibis-server/app/model/metadata/metadata.py +++ b/ibis-server/app/model/metadata/metadata.py @@ -1,5 +1,5 @@ -from app.model.connector import ConnectionInfo -from app.model.metadata.dto import Table, Constraint +from app.model import ConnectionInfo +from app.model.metadata.dto import Constraint, Table class Metadata: diff --git a/ibis-server/app/model/metadata/mysql.py b/ibis-server/app/model/metadata/mysql.py index cad5268ee..b83a7e65a 100644 --- a/ibis-server/app/model/metadata/mysql.py +++ b/ibis-server/app/model/metadata/mysql.py @@ -1,12 +1,13 @@ from json import loads -from app.model.data_source import DataSource, MySqlConnectionInfo +from app.model import MySqlConnectionInfo +from app.model.data_source import DataSource from app.model.metadata.dto import ( - Table, - Constraint, - TableProperties, Column, + Constraint, ConstraintType, + Table, + TableProperties, WrenEngineColumnType, ) from app.model.metadata.metadata import Metadata diff --git a/ibis-server/app/model/metadata/postgres.py b/ibis-server/app/model/metadata/postgres.py index 04013a2d4..50279186b 100644 --- a/ibis-server/app/model/metadata/postgres.py +++ b/ibis-server/app/model/metadata/postgres.py @@ -1,12 +1,13 @@ from json import loads -from app.model.data_source import DataSource, PostgresConnectionInfo +from app.model import PostgresConnectionInfo +from app.model.data_source import DataSource from app.model.metadata.dto import ( - Table, - Constraint, - TableProperties, Column, + Constraint, ConstraintType, + Table, + TableProperties, WrenEngineColumnType, ) from app.model.metadata.metadata import Metadata diff --git a/ibis-server/app/model/validator.py b/ibis-server/app/model/validator.py index 38083fd65..ed6bea04a 100644 --- a/ibis-server/app/model/validator.py +++ b/ibis-server/app/model/validator.py @@ -1,9 +1,6 @@ from __future__ import annotations -from pydantic import BaseModel, Field - from app.model.connector import Connector -from app.model.data_source import ConnectionInfo rules = ["column_is_valid"] @@ -38,12 +35,6 @@ def _validate_column_is_valid(self, parameters: dict[str, str]): raise ValidationError(f"Exception: {type(e)}, message: {str(e)}") -class ValidateDTO(BaseModel): - manifest_str: str = Field(alias="manifestStr", description="Base64 manifest") - parameters: dict[str, str] - connection_info: ConnectionInfo = Field(alias="connectionInfo") - - class ValidationError(Exception): pass diff --git a/ibis-server/app/routers/v2/__init__.py b/ibis-server/app/routers/v2/__init__.py index 2dd7465cc..8828fc479 100644 --- a/ibis-server/app/routers/v2/__init__.py +++ b/ibis-server/app/routers/v2/__init__.py @@ -1,6 +1,6 @@ from fastapi import APIRouter -from app.routers.v2 import ibis, analysis +from app.routers.v2 import analysis, ibis prefix = "/v2" diff --git a/ibis-server/app/routers/v2/analysis.py b/ibis-server/app/routers/v2/analysis.py index 6be53968a..730807b3c 100644 --- a/ibis-server/app/routers/v2/analysis.py +++ b/ibis-server/app/routers/v2/analysis.py @@ -1,17 +1,12 @@ from fastapi import APIRouter -from pydantic import BaseModel, Field from app.logger import log_dto from app.mdl.analyzer import analyze +from app.model import AnalyzeSQLDTO router = APIRouter(prefix="/analysis", tags=["analysis"]) -class AnalyzeSQLDTO(BaseModel): - manifest_str: str = Field(alias="manifestStr", description="Base64 manifest") - sql: str - - @router.get("/sql") @log_dto def analyze_sql(dto: AnalyzeSQLDTO) -> list[dict]: diff --git a/ibis-server/app/routers/v2/ibis.py b/ibis-server/app/routers/v2/ibis.py new file mode 100644 index 000000000..049be04eb --- /dev/null +++ b/ibis-server/app/routers/v2/ibis.py @@ -0,0 +1,57 @@ +from typing import Annotated + +from fastapi import APIRouter, Depends, Query, Response +from fastapi.responses import JSONResponse + +from app.dependencies import verify_query_dto +from app.logger import log_dto +from app.model import ( + QueryDTO, + ValidateDTO, +) +from app.model.connector import ( + Connector, + to_json, +) +from app.model.data_source import DataSource +from app.model.metadata.dto import Constraint, MetadataDTO, Table +from app.model.metadata.factory import MetadataFactory +from app.model.validator import Validator + +router = APIRouter(prefix="/ibis") + + +@router.post("/{data_source}/query", dependencies=[Depends(verify_query_dto)]) +@log_dto +def query( + data_source: DataSource, + dto: QueryDTO, + dry_run: Annotated[bool, Query(alias="dryRun")] = False, +) -> Response: + connector = Connector(data_source, dto.connection_info, dto.manifest_str) + if dry_run: + connector.dry_run(dto.sql) + return Response(status_code=204) + return JSONResponse(to_json(connector.query(dto.sql), dto.column_dtypes)) + + +@router.post("/{data_source}/validate/{rule_name}") +@log_dto +def validate(data_source: DataSource, rule_name: str, dto: ValidateDTO) -> Response: + validator = Validator(Connector(data_source, dto.connection_info, dto.manifest_str)) + validator.validate(rule_name, dto.parameters) + return Response(status_code=204) + + +@router.post("/{data_source}/metadata/tables", response_model=list[Table]) +@log_dto +def get_table_list(data_source: DataSource, dto: MetadataDTO) -> list[Table]: + metadata = MetadataFactory(data_source, dto.connection_info) + return metadata.get_table_list() + + +@router.post("/{data_source}/metadata/constraints", response_model=list[Constraint]) +@log_dto +def get_constraints(data_source: DataSource, dto: MetadataDTO) -> list[Constraint]: + metadata = MetadataFactory(data_source, dto.connection_info) + return metadata.get_constraints() diff --git a/ibis-server/app/routers/v2/ibis/__init__.py b/ibis-server/app/routers/v2/ibis/__init__.py deleted file mode 100644 index 0ba4d36eb..000000000 --- a/ibis-server/app/routers/v2/ibis/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from fastapi import APIRouter - -from app.routers.v2.ibis import bigquery, mysql, postgres, snowflake - -prefix = "/ibis" - -router = APIRouter(prefix=prefix) - -router.include_router(bigquery.router) -router.include_router(mysql.router) -router.include_router(postgres.router) -router.include_router(snowflake.router) diff --git a/ibis-server/app/routers/v2/ibis/bigquery.py b/ibis-server/app/routers/v2/ibis/bigquery.py deleted file mode 100644 index a7252888d..000000000 --- a/ibis-server/app/routers/v2/ibis/bigquery.py +++ /dev/null @@ -1,49 +0,0 @@ -from typing import Annotated - -from fastapi import APIRouter, Query, Response -from fastapi.responses import JSONResponse - -from app.logger import log_dto -from app.model.connector import Connector, QueryBigQueryDTO, to_json -from app.model.data_source import DataSource -from app.model.metadata.dto import MetadataDTO, Table, Constraint -from app.model.metadata.factory import MetadataFactory -from app.model.validator import ValidateDTO, Validator - -router = APIRouter(prefix="/bigquery", tags=["bigquery"]) - -data_source = DataSource.bigquery - - -@router.post("/query") -@log_dto -def query( - dto: QueryBigQueryDTO, dry_run: Annotated[bool, Query(alias="dryRun")] = False -) -> Response: - connector = Connector(data_source, dto.connection_info, dto.manifest_str) - if dry_run: - connector.dry_run(dto.sql) - return Response(status_code=204) - return JSONResponse(to_json(connector.query(dto.sql), dto.column_dtypes)) - - -@router.post("/validate/{rule_name}") -@log_dto -def validate(rule_name: str, dto: ValidateDTO) -> Response: - validator = Validator(Connector(data_source, dto.connection_info, dto.manifest_str)) - validator.validate(rule_name, dto.parameters) - return Response(status_code=204) - - -@router.post("/metadata/tables", response_model=list[Table]) -@log_dto -def get_bigquery_table_list(dto: MetadataDTO) -> list[Table]: - metadata = MetadataFactory(data_source, dto.connection_info) - return metadata.get_table_list() - - -@router.post("/metadata/constraints", response_model=list[Constraint]) -@log_dto -def get_bigquery_constraints(dto: MetadataDTO) -> list[Constraint]: - metadata = MetadataFactory(data_source, dto.connection_info) - return metadata.get_constraints() diff --git a/ibis-server/app/routers/v2/ibis/mysql.py b/ibis-server/app/routers/v2/ibis/mysql.py deleted file mode 100644 index 58919b5a3..000000000 --- a/ibis-server/app/routers/v2/ibis/mysql.py +++ /dev/null @@ -1,50 +0,0 @@ -from typing import Annotated - -from fastapi import APIRouter, Query, Response -from fastapi.responses import JSONResponse - -from app.logger import log_dto -from app.model.connector import Connector, to_json, QueryMySqlDTO -from app.model.data_source import DataSource -from app.model.validator import ValidateDTO, Validator -from app.model.metadata.dto import MetadataDTO, Table, Constraint -from app.model.metadata.factory import MetadataFactory - - -router = APIRouter(prefix="/mysql", tags=["mysql"]) - -data_source = DataSource.mysql - - -@router.post("/query") -@log_dto -def query( - dto: QueryMySqlDTO, dry_run: Annotated[bool, Query(alias="dryRun")] = False -) -> Response: - connector = Connector(data_source, dto.connection_info, dto.manifest_str) - if dry_run: - connector.dry_run(dto.sql) - return Response(status_code=204) - return JSONResponse(to_json(connector.query(dto.sql), dto.column_dtypes)) - - -@router.post("/validate/{rule_name}") -@log_dto -def validate(rule_name: str, dto: ValidateDTO) -> Response: - validator = Validator(Connector(data_source, dto.connection_info, dto.manifest_str)) - validator.validate(rule_name, dto.parameters) - return Response(status_code=204) - - -@router.post("/metadata/tables", response_model=list[Table]) -@log_dto -def get_table_list(dto: MetadataDTO) -> list[Table]: - metadata = MetadataFactory(data_source, dto.connection_info) - return metadata.get_table_list() - - -@router.post("/metadata/constraints", response_model=list[Constraint]) -@log_dto -def get_constraints(dto: MetadataDTO) -> list[Constraint]: - metadata = MetadataFactory(data_source, dto.connection_info) - return metadata.get_constraints() diff --git a/ibis-server/app/routers/v2/ibis/postgres.py b/ibis-server/app/routers/v2/ibis/postgres.py deleted file mode 100644 index dfe31f6bf..000000000 --- a/ibis-server/app/routers/v2/ibis/postgres.py +++ /dev/null @@ -1,49 +0,0 @@ -from typing import Annotated - -from fastapi import APIRouter, Query, Response -from fastapi.responses import JSONResponse - -from app.logger import log_dto -from app.model.connector import Connector, QueryPostgresDTO, to_json -from app.model.data_source import DataSource -from app.model.metadata.dto import MetadataDTO, Table, Constraint -from app.model.metadata.factory import MetadataFactory -from app.model.validator import ValidateDTO, Validator - -router = APIRouter(prefix="/postgres", tags=["postgres"]) - -data_source = DataSource.postgres - - -@router.post("/query") -@log_dto -def query( - dto: QueryPostgresDTO, dry_run: Annotated[bool, Query(alias="dryRun")] = False -) -> Response: - connector = Connector(data_source, dto.connection_info, dto.manifest_str) - if dry_run: - connector.dry_run(dto.sql) - return Response(status_code=204) - return JSONResponse(to_json(connector.query(dto.sql), dto.column_dtypes)) - - -@router.post("/validate/{rule_name}") -@log_dto -def validate(rule_name: str, dto: ValidateDTO) -> Response: - validator = Validator(Connector(data_source, dto.connection_info, dto.manifest_str)) - validator.validate(rule_name, dto.parameters) - return Response(status_code=204) - - -@router.post("/metadata/tables", response_model=list[Table]) -@log_dto -def get_postgres_table_list(dto: MetadataDTO) -> list[Table]: - metadata = MetadataFactory(data_source, dto.connection_info) - return metadata.get_table_list() - - -@router.post("/metadata/constraints", response_model=list[Constraint]) -@log_dto -def get_postgres_constraints(dto: MetadataDTO) -> list[Constraint]: - metadata = MetadataFactory(data_source, dto.connection_info) - return metadata.get_constraints() diff --git a/ibis-server/app/routers/v2/ibis/snowflake.py b/ibis-server/app/routers/v2/ibis/snowflake.py deleted file mode 100644 index 2484584d4..000000000 --- a/ibis-server/app/routers/v2/ibis/snowflake.py +++ /dev/null @@ -1,33 +0,0 @@ -from typing import Annotated - -from fastapi import APIRouter, Query, Response -from fastapi.responses import JSONResponse - -from app.logger import log_dto -from app.model.connector import Connector, QuerySnowflakeDTO, to_json -from app.model.data_source import DataSource -from app.model.validator import ValidateDTO, Validator - -router = APIRouter(prefix="/snowflake", tags=["snowflake"]) - -data_source = DataSource.snowflake - - -@router.post("/query") -@log_dto -def query( - dto: QuerySnowflakeDTO, dry_run: Annotated[bool, Query(alias="dryRun")] = False -) -> Response: - connector = Connector(data_source, dto.connection_info, dto.manifest_str) - if dry_run: - connector.dry_run(dto.sql) - return Response(status_code=204) - return JSONResponse(to_json(connector.query(dto.sql), dto.column_dtypes)) - - -@router.post("/validate/{rule_name}") -@log_dto -def validate(rule_name: str, dto: ValidateDTO) -> Response: - validator = Validator(Connector(data_source, dto.connection_info, dto.manifest_str)) - validator.validate(rule_name, dto.parameters) - return Response(status_code=204) diff --git a/ibis-server/pyproject.toml b/ibis-server/pyproject.toml index a31d3cfb7..2b3f943e2 100644 --- a/ibis-server/pyproject.toml +++ b/ibis-server/pyproject.toml @@ -41,6 +41,9 @@ markers = [ [tool.ruff] target-version = "py311" +[tool.ruff.lint] +select = ["E4", "E7", "E9", "F", "I"] + [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" diff --git a/ibis-server/tests/routers/ibis/test_mysql.py b/ibis-server/tests/routers/ibis/test_mysql.py index 98be9ded2..9b902ef8d 100644 --- a/ibis-server/tests/routers/ibis/test_mysql.py +++ b/ibis-server/tests/routers/ibis/test_mysql.py @@ -3,8 +3,8 @@ import orjson import pytest from fastapi.testclient import TestClient -from testcontainers.mysql import MySqlContainer from sqlalchemy import text +from testcontainers.mysql import MySqlContainer from app.main import app @@ -18,8 +18,8 @@ def file_path(path: str) -> str: return os.path.join(os.path.dirname(__file__), path) - import sqlalchemy import pandas as pd + import sqlalchemy mysql = MySqlContainer("mysql:8.0") mysql.start() diff --git a/ibis-server/tests/routers/ibis/test_postgres.py b/ibis-server/tests/routers/ibis/test_postgres.py index 51f2a4339..7fb981e2f 100644 --- a/ibis-server/tests/routers/ibis/test_postgres.py +++ b/ibis-server/tests/routers/ibis/test_postgres.py @@ -17,8 +17,8 @@ def file_path(path: str) -> str: return os.path.join(os.path.dirname(__file__), path) - import sqlalchemy import pandas as pd + import sqlalchemy pg = PostgresContainer("postgres:16-alpine") pg.start()