teradataml 20.0.0.5__py3-none-any.whl → 20.0.0.6__py3-none-any.whl

teradataml/dataframe/dataframe_utils.py

@@ -1953,7 +1953,7 @@ class DataFrameUtils():
         return _DtypesMappers.DATALAKE_STR_to_TDSQLALCHEMY_DATATYPE_MAPPER[td_str_type]()
 
     @staticmethod
-    def _get_datalake_table_columns_info(schema, table_name, datalake):
+    def _get_datalake_table_columns_info(schema, table_name, datalake, use_dialect=False):
         """
         Function to get column names and corresponding teradatasqlalchemy types
         of a datalake table using results of 'help table <datalake>.<db_name>.<table_name>'
@@ -1990,31 +1990,64 @@ class DataFrameUtils():
               VARCHAR(length=2000, charset='UNICODE'),
               INTEGER()])
         """
-        # Get the column information from the strings type.
-        prepared = preparer(td_dialect())
-        sqlbundle = SQLBundle()
-        full_tbl_name = '{}.{}.{}'.format(prepared.quote(datalake),
-                                          prepared.quote(schema),
-                                          prepared.quote(table_name))
-        help_table_sql = sqlbundle._get_sql_query(SQLConstants.SQL_HELP_TABLE).format(full_tbl_name)
-
-        cur = execute_sql(help_table_sql)
-        td_types_col_index = -1
-        for i, col_metadata in enumerate(cur.description):
-            # Help Table returns column names and
-            # corresponding IcebergType, TeradataInternalType,
-            # TeradataType. We need to extract column index for
-            # 'TeradataType' column.
-            if col_metadata[0].lower() == 'teradatatype':
-                td_types_col_index = i
-
         col_names = []
         col_types = []
-        if td_types_col_index > -1:
-            for col_info in cur.fetchall():
-                col_names.append(col_info[0])
-                col_types.append(DataFrameUtils._get_sqlalchemy_type_from_str(col_info[td_types_col_index]))
+        if not use_dialect:
+            # Get the column information from the strings type.
+            prepared = preparer(td_dialect())
+            sqlbundle = SQLBundle()
+            full_tbl_name = '{}.{}.{}'.format(prepared.quote(datalake),
+                                              prepared.quote(schema),
+                                              prepared.quote(table_name))
+            help_table_sql = sqlbundle._get_sql_query(SQLConstants.SQL_HELP_TABLE).format(full_tbl_name)
+
+            cur = execute_sql(help_table_sql)
+            td_types_col_index = -1
+
+            for i, col_metadata in enumerate(cur.description):
+                # Help Table returns column names and
+                # corresponding IcebergType, TeradataInternalType,
+                # TeradataType. We need to extract column index for
+                # 'TeradataType' column.
+                if col_metadata[0].lower() in ['teradatatype', 'Type']:
+                    td_types_col_index = i
+
+            if td_types_col_index > -1:
+                for col_info in cur.fetchall():
+                    col_names.append(col_info[0])
+                    col_types.append(DataFrameUtils._get_sqlalchemy_type_from_str(col_info[td_types_col_index]))
+            else:
+                raise TeradataMlException(Messages.get_message(MessageCodes.TDMLDF_CREATE_FAIL),
+                                          MessageCodes.TDMLDF_CREATE_FAIL)
         else:
-            raise TeradataMlException(Messages.get_message(MessageCodes.TDMLDF_CREATE_FAIL),
-                                      MessageCodes.TDMLDF_CREATE_FAIL)
+            new_kwarg = get_connection().dialect.__class__.__name__ + "_datalake"
+            all_col_info = get_connection().dialect.get_columns(connection=get_connection(),
+                                                                table_name=table_name,
+                                                                schema=schema,
+                                                                table_only=True,
+                                                                **{new_kwarg: datalake})
+            for col_dict in all_col_info:
+                col_names.append(col_dict.get('name', col_dict.get('Column Name')))
+                col_types.append(col_dict.get('type', col_dict.get('Type')))
+
         return col_names, col_types
+
+    @staticmethod
+    def check_otf_dataframe():
+        """Decorator for validating if DataFrame is created on OTF table or not and throw error."""
+        def decorator(method):
+            def wrapper(self, *args, **kwargs):
+                if not self._datalake:
+                    attr = getattr(type(self), method.__name__, None)
+                    caller_name = method.__name__ + '()'
+                    if isinstance(attr, property):
+                        caller_name = method.__name__
+                    raise TeradataMlException(Messages.get_message(MessageCodes.OTF_TABLE_REQUIRED,
+                                                                   caller_name),
+                                              MessageCodes.UNSUPPORTED_OPERATION)
+
+                return method(self, *args, **kwargs)
+
+            return wrapper
+
+        return decorator