feat(db engines): add support for Opendistro Elasticsearch (AWS ES) (#12602)

* feat(db engines): add support for Opendistro Elasticsearch (AWS ES)

* add time grains

* lint

* bump elasticsearch-dbapi version

* add tests

* fix test

setup.py

@@ -126,7 +126,7 @@ setup(
         "drill": ["sqlalchemy-drill==0.1.dev"],
         "druid": ["pydruid>=0.6.1,<0.7"],
         "solr": ["sqlalchemy-solr >= 0.2.0"],
-        "elasticsearch": ["elasticsearch-dbapi>=0.1.0, <0.2.0"],
+        "elasticsearch": ["elasticsearch-dbapi>=0.2.0, <0.3.0"],
         "exasol": ["sqlalchemy-exasol>=2.1.0, <2.2"],
         "excel": ["xlrd>=1.2.0, <1.3"],
         "gsheets": ["gsheetsdb>=0.1.9"],

superset/db_engine_specs/elasticsearch.py

@@ -61,3 +61,34 @@ class ElasticSearchEngineSpec(BaseEngineSpec):  # pylint: disable=abstract-metho
         if target_type.upper() == utils.TemporalType.DATETIME:
             return f"""CAST('{dttm.isoformat(timespec="seconds")}' AS DATETIME)"""
         return None
+
+
+class OpenDistroEngineSpec(BaseEngineSpec):  # pylint: disable=abstract-method
+
+    time_groupby_inline = True
+    time_secondary_columns = True
+    allows_joins = False
+    allows_subqueries = True
+
+    _time_grain_expressions = {
+        None: "{col}",
+        "PT1S": "date_format({col}, 'yyyy-MM-dd HH:mm:ss.000')",
+        "PT1M": "date_format({col}, 'yyyy-MM-dd HH:mm:00.000')",
+        "PT1H": "date_format({col}, 'yyyy-MM-dd HH:00:00.000')",
+        "P1D": "date_format({col}, 'yyyy-MM-dd 00:00:00.000')",
+        "P1M": "date_format({col}, 'yyyy-MM-01 00:00:00.000')",
+        "P1Y": "date_format({col}, 'yyyy-01-01 00:00:00.000')",
+    }
+
+    engine = "odelasticsearch"
+    engine_name = "ElasticSearch"
+
+    @classmethod
+    def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]:
+        if target_type.upper() == utils.TemporalType.DATETIME:
+            return f"""'{dttm.isoformat(timespec="seconds")}'"""
+        return None
+
+    @staticmethod
+    def _mutate_label(label: str) -> str:
+        return label.replace(".", "_")

tests/db_engine_specs/elasticsearch_tests.py

@@ -14,7 +14,12 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from superset.db_engine_specs.elasticsearch import ElasticSearchEngineSpec
+from sqlalchemy import column
+
+from superset.db_engine_specs.elasticsearch import (
+    ElasticSearchEngineSpec,
+    OpenDistroEngineSpec,
+)
 from tests.db_engine_specs.base_tests import TestDbEngineSpec
 
 
@@ -26,3 +31,26 @@ class TestElasticSearchDbEngineSpec(TestDbEngineSpec):
             ElasticSearchEngineSpec.convert_dttm("DATETIME", dttm),
             "CAST('2019-01-02T03:04:05' AS DATETIME)",
         )
+
+    def test_opendistro_convert_dttm(self):
+        """
+        DB Eng Specs (opendistro): Test convert_dttm
+        """
+        dttm = self.get_dttm()
+
+        self.assertEqual(
+            OpenDistroEngineSpec.convert_dttm("DATETIME", dttm),
+            "'2019-01-02T03:04:05'",
+        )
+
+    def test_opendistro_sqla_column_label(self):
+        """
+        DB Eng Specs (opendistro): Test column label
+        """
+        test_cases = {
+            "Col": "Col",
+            "Col.keyword": "Col_keyword",
+        }
+        for original, expected in test_cases.items():
+            actual = OpenDistroEngineSpec.make_label_compatible(column(original).name)
+            self.assertEqual(actual, expected)