feat: add connector for CouchbaseDB

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)]
+    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,
+        }