From a8085d318c7fdd4b4591c45dc8eb6f399cabb670 Mon Sep 17 00:00:00 2001
From: "Piyush:)" <47020544+piyushsarraf@users.noreply.github.com>
Date: Wed, 19 Apr 2023 12:58:20 +0530
Subject: [PATCH] feat: Adding Custom-Query support for DB2.  (#807)

* feat: Adding Custom-Query Support for DB2

* fix: typo in README.md

---------

Co-authored-by: Piyush Sarraf <piyushsarraf@google.com>
---
 README.md                           |  2 +-
 third_party/ibis/ibis_DB2/client.py | 43 +++++++++++++++++++++++++++--
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 06018a9d6..0987b817a 100644
--- a/README.md
+++ b/README.md
@@ -233,7 +233,7 @@ data-validation (--verbose or -v) (--log-level or -ll) generate-table-partitions
                         Local: Provide a relative path of the target directory. Eg: `partitions_dir`
   --partition-num [1-1000], -pn [1-1000]
                         Number of partitions/config files to generate
-                        In case this value exceeds the row count of the source/target table, its will be decreased to max(source_row_count, target_row_count)
+                        In case this value exceeds the row count of the source/target table, it will be decreased to max(source_row_count, target_row_count)
   [--partition-key PARTITION_KEY, -partkey PARTITION_KEY]
                         Column on which the partitions would be generated. Column type must be integer. Defaults to Primary key
   [--filters SOURCE_FILTER:TARGET_FILTER]
diff --git a/third_party/ibis/ibis_DB2/client.py b/third_party/ibis/ibis_DB2/client.py
index cd2d12e64..2438a506a 100644
--- a/third_party/ibis/ibis_DB2/client.py
+++ b/third_party/ibis/ibis_DB2/client.py
@@ -17,9 +17,10 @@
 from typing import Optional
 
 import sqlalchemy as sa
-
+import ibis.expr.operations as ops
 import third_party.ibis.ibis_DB2.alchemy as alch
 from third_party.ibis.ibis_DB2.compiler import DB2Dialect
+import ibis.expr.schema as sch
 
 
 class DB2Table(alch.AlchemyTable):
@@ -198,4 +199,42 @@ def list_tables(self, like=None, database=None, schema=None):
             return parent.list_tables(like=like, schema=schema)
 
     def get_schema(self, name, schema=None):
-        return self.table(name, schema=schema).schema()
\ No newline at end of file
+        return self.table(name, schema=schema).schema()
+
+    def sql(self, query):
+        """
+        Convert a DB2 query to an Ibis table expression
+
+        Parameters
+        ----------
+        query: string
+           SQL query to execute on connection
+
+        Returns
+        -------
+        table : TableExpr
+        """
+        limited_query = 'SELECT * FROM ({}) t0 LIMIT 1'.format(query)
+        schema = self._get_schema_using_query(limited_query)
+        return ops.SQLQueryResult(query, schema, self).to_expr()
+
+    def _get_schema_using_query(self, limited_query):
+        type_map = {
+            'INTEGER': 'int64',
+            'BOOLEAN': 'boolean',
+            'FLOAT': 'float64',
+            'VARCHAR': 'string',
+            'DOUBLE': 'float64'
+        }
+
+        with self._execute(limited_query, results=True) as cur:
+            names = []
+            ibis_types = []
+            for row in cur.proxy._cursor_description():
+                names.append(row[0].lower())
+                datatypes = list(row[1])
+                for datatype in datatypes: 
+                    if datatype in type_map:
+                        ibis_types.append(type_map[datatype])
+                        break
+        return sch.Schema(names, ibis_types)
\ No newline at end of file