diff --git a/superset/db_engine_specs/couchbase.py b/superset/db_engine_specs/couchbase.py index 403b2fb33e1e2..a5576719568c3 100644 --- a/superset/db_engine_specs/couchbase.py +++ b/superset/db_engine_specs/couchbase.py @@ -14,22 +14,52 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from superset.db_engine_specs.base import BaseEngineSpec -from superset.constants import TimeGrain -from datetime import datetime -import sqlparse -from superset.db_engine_specs.base import BaseEngineSpec -from superset.superset_typing import ResultSetColumnType, SQLAColumnType +# pylint: disable=too-many-lines +from __future__ import annotations + +from sqlalchemy.engine.url import URL from flask_babel import gettext as __ -from superset.db_engine_specs.base import BaseEngineSpec +from superset.superset_typing import ResultSetColumnType +from marshmallow import fields, Schema +from marshmallow.validate import Range +from sqlalchemy import column +from superset.databases.utils import make_url_safe +from superset.db_engine_specs.base import BaseEngineSpec, BasicParametersMixin +from superset.constants import TimeGrain +import os +import datetime from typing import ( - Any + Any, + TypedDict ) +class BasicParametersType(TypedDict, total=False): + username: str | None + password: str | None + host: str + port: int + database: str + query: dict[str, Any] + encryption: bool + +class CouchbaseParametersSchema(Schema): + username = fields.String(allow_none=True, metadata={"description": __("Username")}) + password = fields.String(allow_none=True, metadata={"description": __("Password")}) + host = fields.String(required=True, metadata={"description": __("Hostname or IP address")}) + port = fields.Integer(allow_none=True, metadata={"description": __("Database port")}, validate=Range(min=0, max=65535)) + database = fields.String(allow_none=True, metadata={"description": __("Database name")}) + encryption = fields.Boolean(dump_default=False, metadata={"description": __("Use an encrypted connection to the database")}) + query = fields.Dict(keys=fields.Str(), values=fields.Raw(), metadata={"description": __("Additional parameters")}) -class CouchbaseEngineSpec(BaseEngineSpec): +class CouchbaseEngineSpec(BasicParametersMixin,BaseEngineSpec): engine = 'couchbase' - engine_name = 'couchbase' + engine_name = 'Couchbase Columnar' + allows_joins = False allows_subqueries = False + default_driver = 'couchbase' + sqlalchemy_uri_placeholder = ("columnar+couchbase://user:password@host[:port][/dbname][?ssl=value&=value...]") + parameters_schema = CouchbaseParametersSchema() + encryption_parameters = {"sslmode": "require"} + _time_grain_expressions = { None: "{col}", TimeGrain.SECOND: "DATE_TRUNC_STR(TOSTRING({col}),'second')", @@ -38,16 +68,16 @@ class CouchbaseEngineSpec(BaseEngineSpec): TimeGrain.DAY: "DATE_TRUNC_STR(TOSTRING({col}),'day')", TimeGrain.MONTH: "DATE_TRUNC_STR(TOSTRING({col}),'month')", TimeGrain.YEAR: "DATE_TRUNC_STR(TOSTRING({col}),'year')", - TimeGrain.QUARTER: "DATE_TRUNC_STR(TOSTRING({col}),'quarter')" + TimeGrain.QUARTER: "DATE_TRUNC_STR(TOSTRING({col}),'quarter')" } @classmethod def epoch_to_dttm(cls) -> str: - return "{col} * 1000" + return "MILLIS_TO_STR({col} * 1000, '111')" @classmethod def epoch_ms_to_dttm(cls) -> str: - return "{col}" + return "MILLIS_TO_STR({col}, '111')" @classmethod def convert_dttm( @@ -57,6 +87,42 @@ def convert_dttm( return f"DATETIME(DATE_FORMAT_STR(STR_TO_UTC('{dttm.date().isoformat()}'), 'iso8601'))" @classmethod - def parse_sql(cls, sql: str) -> list[str]: - sql.replace("`COUNT(*)`","COUNT(*)") - return [str(s).strip(" ;") for s in sqlparse.parse(sql)] \ No newline at end of file + def build_sqlalchemy_uri(cls, parameters: dict, encrypted_extra=None): + query_params = parameters.get("query", {}).copy() + if parameters.get("encryption", False): + query_params.update(cls.encryption_parameters) + + uri = URL.create( + f"{cls.engine}+{cls.default_driver}", + username=parameters.get("username"), + password=parameters.get("password"), + host=parameters["host"], + port=parameters.get("port", 18091), # Default SSL port for Couchbase + database=parameters.get("database", "default"), + query=query_params + ) + return str(uri) + + + @classmethod + def get_parameters_from_uri( # pylint: disable=unused-argument + cls, uri: str, encrypted_extra: dict[str, Any] | None = None + ) -> BasicParametersType: + url = make_url_safe(uri) + query = { + key: value + for (key, value) in url.query.items() + if (key, value) not in cls.encryption_parameters.items() + } + encryption = all( + item in url.query.items() for item in cls.encryption_parameters.items() + ) + return { + "username": url.username, + "password": url.password, + "host": url.host, + "port": url.port, + "database": url.database, + "query": query, + "encryption": encryption, + } \ No newline at end of file diff --git a/tests/unit_tests/db_engine_specs/test_couchbase.py b/tests/unit_tests/db_engine_specs/test_couchbase.py new file mode 100644 index 0000000000000..e69de29bb2d1d