feat: Issue356 db2 test

data_validation/cli_tools.py

@@ -125,6 +125,15 @@
             "Desired Kerberos service name ('impala' if not provided)",
         ],
     ],
+    "DB2": [
+        ["host", "Desired DB2 host"],
+        ["port", "Desired DB2 port (50000 if not provided)"],
+        ["user", "Username to connect to"],
+        ["password", "Password for authentication of user"],
+        ["database", "Database in DB2 to connect to"],
+        ["url", "URL link in DB2 to connect to"],
+        ["driver", "Driver link in DB2 to connect to (default ibm_db_sa)"],
+    ],
 }
 
 

data_validation/clients.py

@@ -30,6 +30,7 @@
 import third_party.ibis.ibis_addon.datatypes
 from third_party.ibis.ibis_cloud_spanner.api import connect as spanner_connect
 from third_party.ibis.ibis_impala.api import impala_connect
+
 from data_validation import client_info
 from data_validation import consts, exceptions
 
@@ -83,6 +84,12 @@ def get_client_call(*args, **kwargs):
         "pip install snowflake-connector-python"
     )
 
+# If you have Db2 client installed
+try:
+    from third_party.ibis.ibis_DB2.client import DB2Client
+except Exception:
+    DB2Client = _raise_missing_client_error("pip install ibm_db_sa")
+
 
 def get_bigquery_client(project_id, dataset_id=None, credentials=None):
     info = client_info.get_http_client_info()
@@ -129,6 +136,7 @@ def get_ibis_table(client, schema_name, table_name, database_name=None):
     if type(client) in [
         OracleClient,
         PostgreSQLClient,
+        DB2Client,
         MSSQLClient,
     ]:
         return client.table(table_name, database=database_name, schema=schema_name)
@@ -143,6 +151,7 @@ def list_schemas(client):
     if type(client) in [
         OracleClient,
         PostgreSQLClient,
+        DB2Client,
         MSSQLClient,
     ]:
         return client.list_schemas()
@@ -157,6 +166,7 @@ def list_tables(client, schema_name):
     if type(client) in [
         OracleClient,
         PostgreSQLClient,
+        DB2Client,
         MSSQLClient,
     ]:
         return client.list_tables(schema=schema_name)
@@ -235,4 +245,5 @@ def get_data_client(connection_config):
     "MSSQL": MSSQLClient,
     "Snowflake": snowflake_connect,
     "Spanner": spanner_connect,
+    "DB2": DB2Client,
 }

data_validation/consts.py

@@ -125,7 +125,16 @@
 RESULT_TYPE_TARGET = "target"
 
 # Ibis Object Info
-NUMERIC_DATA_TYPES = ["float64", "int8", "int16", "int32", "int64", "decimal"]
+NUMERIC_DATA_TYPES = [
+    "float64",
+    "int32",
+    "int64",
+    "decimal",
+    "int32[non-nullable]",
+    "int64[non-nullable]",
+    "float64[non-nullable]",
+    "decimal[non-nullable]",
+]
 
 FORMAT_TYPES = ["csv", "json", "table", "text"]
 

docs/connections.md

@@ -53,6 +53,7 @@ The data validation tool supports the following connection types.
 * [FileSystem](#filesystem)
 * [Impala](#Impala)
 * [Hive](#Hive)
+* [DB2] (#DB2)
 
 As you see above, Teradata and BigQuery have different sets of custom arguments (for example project_id for BQ versus host for Teradata).
 
@@ -285,3 +286,19 @@ Please note that for Group By validations, the following property must be set in
     "auth_mechanism":"PLAIN"
 }
 ```
+## DB2
+```
+{
+    # Configuration Required for All Data Sources
+    "source_type": "DB2",
+
+    # Connection Details
+    "host": "localhost",
+    "port": 50000,
+    "driver": "ibm_db_sa",
+    "user": "my-username",
+    "password": "my-password",
+    "database": "my-db",
+    "url": "my-url",
+}
+```

third_party/ibis/ibis_DB2/alchemy.py

@@ -18,7 +18,7 @@
 
 import ibis.expr.datatypes as dt11
 import ibis.expr.datatypes as dts
-import ibis.sql.alchemy as s_al
+import ibis.backends.base_sqlalchemy.alchemy as s_al
 import third_party.ibis.ibis_DB2.expr.datatypes as dt
 
 _ibis_type_to_sqla = {

third_party/ibis/ibis_DB2/api.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ibis.sql.alchemy import to_sqlalchemy
+from ibis.backends.base_sqlalchemy.alchemy import to_sqlalchemy
 from third_party.ibis.ibis_DB2.client import DB2Client
 from third_party.ibis.ibis_DB2.compiler import dialect, rewrites  # noqa: F401
 

third_party/ibis/ibis_DB2/compiler.py

@@ -31,11 +31,11 @@
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
 import ibis.expr.types as ir
-import ibis.sql.alchemy as alch
+import ibis.backends.base_sqlalchemy.alchemy as alch
 import third_party.ibis.ibis_DB2.alchemy as db2_alch
 
 # used for literal translate
-from ibis.sql.alchemy import (
+from ibis.backends.base_sqlalchemy.alchemy import (
     _get_sqla_table,
     _variance_reduction,
     fixed_arity,
@@ -363,6 +363,29 @@ def reduction_compiler(t, expr):
 
     return reduction_compiler
 
+def _count_start(sa_func):
+    return sa_func
+
+def _reduction_count(sa_func):
+    def formatter(t, expr):
+        op = expr.op()
+        *args, where = op.args
+
+        return _reduction_format(t, sa_func, where, *args)
+
+    return formatter
+
+def _reduction_format(t, sa_func, where, arg, *args):
+    if where is not None:
+        arg = t.translate(where.ifelse(arg, None))
+    else:
+        arg = t.translate(arg)
+
+    #Db2 doesn't allow '*' to be parameterized, probably better way to fix this... 
+    if arg == '*':
+        arg = None
+
+    return sa_func(arg, *map(t.translate, args))
 
 def _log(t, expr):
     arg, base = expr.op().args
@@ -572,6 +595,9 @@ def _day_of_week_name(t, expr):
         ops.CumulativeAny: unary(sa.func.bool_or),
         ops.IdenticalTo: _identical_to,
         ops.HLLCardinality: _hll_cardinality,
+        # aggregate methods
+        ops.Count: _reduction_count(sa.func.count),
+
     }
 )