Skip to content

Mask sensitive information #672

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 16, 2024
Merged
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
72 changes: 36 additions & 36 deletions ibis-server/app/model/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, SecretStr

manifest_str_field = Field(alias="manifestStr", description="Base64 manifest")
connection_info_field = Field(alias="connectionInfo")
Expand Down Expand Up @@ -46,70 +46,70 @@ class QueryTrinoDTO(QueryDTO):


class BigQueryConnectionInfo(BaseModel):
project_id: str
dataset_id: str
credentials: str = Field(description="Base64 encode `credentials.json`")
project_id: SecretStr
dataset_id: SecretStr
credentials: SecretStr = Field(description="Base64 encode `credentials.json`")


class ClickHouseConnectionInfo(BaseModel):
host: str
port: int
database: str
user: str
password: str
host: SecretStr
port: SecretStr
database: SecretStr
user: SecretStr
password: SecretStr


class MSSqlConnectionInfo(BaseModel):
host: str
port: int
database: str
user: str
password: str
host: SecretStr
port: SecretStr
database: SecretStr
user: SecretStr
password: SecretStr
driver: str = Field(
default="FreeTDS",
description="On Mac and Linux this is usually `FreeTDS. On Windows, it is usually `ODBC Driver 18 for SQL Server`",
)


class MySqlConnectionInfo(BaseModel):
host: str
port: int
database: str
user: str
password: str
host: SecretStr
port: SecretStr
database: SecretStr
user: SecretStr
password: SecretStr


class ConnectionUrl(BaseModel):
connection_url: str = Field(alias="connectionUrl")
connection_url: SecretStr = Field(alias="connectionUrl")


class PostgresConnectionInfo(BaseModel):
host: str = Field(examples=["localhost"])
port: int = Field(examples=[5432])
database: str
user: str
password: str
host: SecretStr = Field(examples=["localhost"])
port: SecretStr = Field(examples=[5432])
database: SecretStr
user: SecretStr
password: SecretStr


class SnowflakeConnectionInfo(BaseModel):
user: str
password: str
account: str
database: str
sf_schema: str = Field(
user: SecretStr
password: SecretStr
account: SecretStr
database: SecretStr
sf_schema: SecretStr = Field(
alias="schema"
) # Use `sf_schema` to avoid `schema` shadowing in BaseModel


class TrinoConnectionInfo(BaseModel):
host: str
port: int = Field(default=8080)
catalog: str
trino_schema: str = Field(
host: SecretStr
port: SecretStr = Field(default="8080")
catalog: SecretStr
trino_schema: SecretStr = Field(
alias="schema"
) # Use `trino_schema` to avoid `schema` shadowing in BaseModel
user: str | None = None
password: str | None = None
user: SecretStr | None = None
password: SecretStr | None = None


ConnectionInfo = (
Expand Down
76 changes: 43 additions & 33 deletions ibis-server/app/model/data_source.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,87 +72,97 @@ def get_connection(self, info: ConnectionInfo) -> BaseBackend:

@staticmethod
def get_bigquery_connection(info: BigQueryConnectionInfo) -> BaseBackend:
credits_json = loads(base64.b64decode(info.credentials).decode("utf-8"))
credits_json = loads(
base64.b64decode(info.credentials.get_secret_value()).decode("utf-8")
)
credentials = service_account.Credentials.from_service_account_info(
credits_json
)
return ibis.bigquery.connect(
project_id=info.project_id,
dataset_id=info.dataset_id,
project_id=info.project_id.get_secret_value(),
dataset_id=info.dataset_id.get_secret_value(),
credentials=credentials,
)

@staticmethod
def get_clickhouse_connection(
info: ConnectionUrl | ClickHouseConnectionInfo,
) -> BaseBackend:
connection_url = (
getattr(info, "connection_url", None)
or f"clickhouse://{info.user}:{info.password}@{info.host}:{info.port}/{info.database}"
)
return ibis.connect(
connection_url,
if hasattr(info, "connection_url"):
url = info.connection_url.get_secret_value()
# ibis miss port of connection url, so we need to pass it explicitly
port=urlparse(connection_url).port,
return ibis.connect(url, port=urlparse(url).port)
return ibis.clickhouse.connect(
host=info.host.get_secret_value(),
port=int(info.port.get_secret_value()),
database=info.database.get_secret_value(),
user=info.user.get_secret_value(),
password=info.password.get_secret_value(),
)

@staticmethod
def get_mssql_connection(info: MSSqlConnectionInfo) -> BaseBackend:
# mssql in ibis does not support connection url
return ibis.mssql.connect(
host=info.host,
port=info.port,
database=info.database,
user=info.user,
password=info.password,
host=info.host.get_secret_value(),
port=info.port.get_secret_value(),
database=info.database.get_secret_value(),
user=info.user.get_secret_value(),
password=info.password.get_secret_value(),
driver=info.driver,
)

@staticmethod
def get_mysql_connection(
info: ConnectionUrl | MySqlConnectionInfo,
) -> BaseBackend:
connection_url = (
getattr(info, "connection_url", None)
or f"mysql://{info.user}:{info.password}@{info.host}:{info.port}/{info.database}"
)
return ibis.connect(
connection_url,
if hasattr(info, "connection_url"):
url = info.connection_url.get_secret_value()
# ibis miss port of connection url, so we need to pass it explicitly
port=urlparse(connection_url).port,
return ibis.connect(url, port=urlparse(url).port)
return ibis.mysql.connect(
host=info.host.get_secret_value(),
port=int(info.port.get_secret_value()),
database=info.database.get_secret_value(),
user=info.user.get_secret_value(),
password=info.password.get_secret_value(),
)

@staticmethod
def get_postgres_connection(
info: ConnectionUrl | PostgresConnectionInfo,
) -> BaseBackend:
return ibis.connect(
getattr(info, "connection_url", None)
or f"postgres://{info.user}:{info.password}@{info.host}:{info.port}/{info.database}"
if hasattr(info, "connection_url"):
return ibis.connect(info.connection_url.get_secret_value())
return ibis.postgres.connect(
host=info.host.get_secret_value(),
port=int(info.port.get_secret_value()),
database=info.database.get_secret_value(),
user=info.user.get_secret_value(),
password=info.password.get_secret_value(),
)

@staticmethod
def get_snowflake_connection(info: SnowflakeConnectionInfo) -> BaseBackend:
return ibis.snowflake.connect(
user=info.user,
password=info.password,
account=info.account,
database=info.database,
schema=info.sf_schema,
user=info.user.get_secret_value(),
password=info.password.get_secret_value(),
account=info.account.get_secret_value(),
database=info.database.get_secret_value(),
schema=info.sf_schema.get_secret_value(),
)

@staticmethod
def get_trino_connection(
info: ConnectionUrl | TrinoConnectionInfo,
) -> BaseBackend:
if hasattr(info, "connection_url"):
return ibis.connect(info.connection_url)

return ibis.connect(info.connection_url.get_secret_value())
return ibis.trino.connect(
host=info.host,
port=info.port,
database=info.catalog,
schema=info.trino_schema,
user=info.user,
password=info.password,
password=info.password.get_secret_value(),
)
Loading