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 <[email protected]>
GoogleCloudPlatform · Apr 19, 2023 · a8085d3 · a8085d3
1 parent c16e2f7
commit a8085d3
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 3 deletions.
diff --git 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
@@ -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()
+        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)