Add integration tests for Apache Pinot (apache#13195)

* Add integration tests for Apache Pinot * fixup! Add integration tests for Apache Pinot * fixup! fixup! Add integration tests for Apache Pinot * fixup! fixup! fixup! Add integration tests for Apache Pinot * fixup! fixup! fixup! fixup! Add integration tests for Apache Pinot * Update setup.cfg
aniejo · Dec 30, 2020 · 98f097e · 98f097e
1 parent 57cbcf6
commit 98f097e
Show file tree

Hide file tree

Showing 9 changed files with 73 additions and 5 deletions.
diff --git a/BREEZE.rst b/BREEZE.rst
@@ -2301,7 +2301,7 @@ This is the current syntax for  `./breeze <./breeze>`_:
           start all integrations. Selected integrations are not saved for future execution.
           One of:
 
-                 cassandra kerberos mongo openldap presto rabbitmq redis all
+                 cassandra kerberos mongo openldap pinot presto rabbitmq redis all
 
   --init-script INIT_SCRIPT_FILE
           Initialization script name - Sourced from files/airflow-breeze-config. Default value

diff --git a/TESTING.rst b/TESTING.rst
@@ -279,6 +279,8 @@ The following integrations are available:
      - Integration required for MongoDB hooks
    * - openldap
      - Integration required for OpenLDAP hooks
+   * - pinot
+     - Integration required for Apache Pinot hooks
    * - presto
      - Integration required for Presto hooks
    * - rabbitmq

diff --git a/breeze-complete b/breeze-complete
@@ -25,7 +25,7 @@
 
 _breeze_allowed_python_major_minor_versions="2.7 3.5 3.6 3.7 3.8"
 _breeze_allowed_backends="sqlite mysql postgres"
-_breeze_allowed_integrations="cassandra kerberos mongo openldap presto rabbitmq redis all"
+_breeze_allowed_integrations="cassandra kerberos mongo openldap pinot presto rabbitmq redis all"
 _breeze_allowed_kubernetes_modes="image"
 _breeze_allowed_kubernetes_versions="v1.18.6 v1.17.5 v1.16.9"
 _breeze_allowed_helm_versions="v3.2.4"

diff --git a/scripts/ci/docker-compose/integration-pinot.yml b/scripts/ci/docker-compose/integration-pinot.yml
@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+---
+version: "2.2"
+services:
+  pinot:
+    image: apachepinot/pinot:latest
+    ports:
+      - "9080:9080"
+    volumes:
+      - /dev/urandom:/dev/random   # Required to get non-blocking entropy source
+    command: QuickStart -type batch
+
+  airflow:
+    environment:
+      - INTEGRATION_PINOT=true
+    depends_on:
+      - pinot
diff --git a/scripts/ci/libraries/_initialization.sh b/scripts/ci/libraries/_initialization.sh
@@ -225,7 +225,7 @@ function initialization::initialize_dockerhub_variables() {
 
 # Determine available integrations
 function initialization::initialize_available_integrations() {
-    export AVAILABLE_INTEGRATIONS="cassandra kerberos mongo openldap presto rabbitmq redis"
+    export AVAILABLE_INTEGRATIONS="cassandra kerberos mongo openldap pinot presto rabbitmq redis"
 }
 
 # Needs to be declared outside of function for MacOS

diff --git a/scripts/in_container/check_environment.sh b/scripts/in_container/check_environment.sh
@@ -147,6 +147,17 @@ check_integration "Presto (HTTP)" "presto" "nc -zvv presto 8080" 40
 check_integration "Presto (HTTPS)" "presto" "nc -zvv presto 7778" 40
 check_integration "Presto (API)" "presto" \
     "curl --max-time 1 http://presto:8080/v1/info/ | grep '\"starting\":false'" 20
+check_integration "Pinot (HTTP)" "pinot" "nc -zvv pinot 9000" 40
+check_integration "Presto (Controller API)" "pinot" \
+    "curl --max-time 1 -X GET 'http://pinot:9000/health' " \
+    "-H 'accept: text/plain' | grep OK" 20
+check_integration "Presto (Controller API)" "pinot" \
+    "curl --max-time 1 -X GET 'http://pinot:9000/pinot-controller/admin' " \
+    "-H 'accept: text/plain' | grep GOOD" 20
+check_integration "Presto (Broker API)" "pinot" \
+    "curl --max-time 1 -X GET 'http://pinot:8000/health' " \
+    "-H 'accept: text/plain' | grep OK" 20
+
 echo "-----------------------------------------------------------------------------------------------"
 
 if [[ ${EXIT_CODE} != 0 ]]; then

diff --git a/setup.cfg b/setup.cfg
@@ -113,7 +113,12 @@ install_requires =
     pendulum~=2.0
     pep562~=1.0;python_version<"3.7"
     psutil>=4.2.0, <6.0.0
+    # snowflake-connector-python don't support newer version
+    # https://github.com/snowflakedb/snowflake-connector-python/blob/v2.3.6/setup.py#L201
+    pyOpenSSL<20.0.0
     pygments>=2.0.1, <3.0
+    # snowflake-connector-python, flask-jwt-extended, msal are not compatible with newest version.
+    pyjwt<2
     python-daemon>=2.1.1
     python-dateutil>=2.3, <3
     python-nvd3~=0.15.0

diff --git a/setup.py b/setup.py
@@ -345,7 +345,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version
     'flask-bcrypt>=0.7.1',
 ]
 pinot = [
-    'pinotdb==0.1.1',
+    'pinotdb>=0.1.1,<1.0.0',
 ]
 plexus = [
     'arrow>=0.16.0',
@@ -397,7 +397,10 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version
     # This library monkey patches the requests library, so SSL is broken globally.
     # See: https://github.com/snowflakedb/snowflake-connector-python/issues/324
     'requests<2.24.0',
-    'snowflake-connector-python>=1.5.2',
+    # Newest version drop support for old version of azure-storage-blob
+    # Until #12188 is solved at least we need to limit maximum version.
+    # https://github.com/apache/airflow/pull/12188
+    'snowflake-connector-python>=1.5.2,<=2.3.6',
     'snowflake-sqlalchemy>=1.1.0',
 ]
 spark = [
@@ -460,6 +463,9 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version
     'importlib-resources~=1.4',
     'ipdb',
     'jira',
+    # HACK: Moto is not compatible with newer versions
+    # See: https://github.com/spulec/moto/issues/3535
+    'mock<4.0.3',
     'mongomock',
     'moto',
     'mypy==0.770',

diff --git a/tests/providers/apache/pinot/hooks/test_pinot.py b/tests/providers/apache/pinot/hooks/test_pinot.py
@@ -23,6 +23,8 @@
 import unittest
 from unittest import mock
 
+import pytest
+
 from airflow.exceptions import AirflowException
 from airflow.providers.apache.pinot.hooks.pinot import PinotAdminHook, PinotDbApiHook
 
@@ -264,3 +266,13 @@ def test_get_pandas_df(self):
         self.assertEqual(column, df.columns[0])
         for i in range(len(result_sets)):  # pylint: disable=consider-using-enumerate
             self.assertEqual(result_sets[i][0], df.values.tolist()[i][0])
+
+
+class TestPinotDbApiHookIntegration(unittest.TestCase):
+    @pytest.mark.integration("pinot")
+    @mock.patch.dict('os.environ', AIRFLOW_CONN_PINOT_BROKER_DEFAULT="pinot://pinot:8000/")
+    def test_should_return_records(self):
+        hook = PinotDbApiHook()
+        sql = "select playerName from baseballStats  ORDER BY playerName limit 5"
+        records = hook.get_records(sql)
+        self.assertEqual([["A. Harry"], ["A. Harry"], ["Aaron"], ["Aaron Albert"], ["Aaron Albert"]], records)