allows naming conventions to be changed

dlt/destinations/impl/bigquery/sql_client.py

@@ -1,5 +1,5 @@
 from contextlib import contextmanager
-from typing import Any, AnyStr, ClassVar, Iterator, List, Optional, Sequence
+from typing import Any, AnyStr, ClassVar, Iterator, List, Optional, Sequence, Generator
 
 import google.cloud.bigquery as bigquery  # noqa: I250
 from google.api_core import exceptions as api_core_exceptions
@@ -8,6 +8,7 @@
 from google.cloud.bigquery.dbapi import Connection as DbApiConnection, Cursor as BQDbApiCursor
 from google.cloud.bigquery.dbapi import exceptions as dbapi_exceptions
 
+from dlt.common import logger
 from dlt.common.configuration.specs import GcpServiceAccountCredentialsWithoutDefaults
 from dlt.common.destination import DestinationCapabilitiesContext
 from dlt.common.typing import StrAny
@@ -44,17 +45,30 @@ class BigQueryDBApiCursorImpl(DBApiCursorImpl):
     """Use native BigQuery data frame support if available"""
 
     native_cursor: BQDbApiCursor  # type: ignore
+    df_iterator: Generator[Any, None, None]
 
-    def df(self, chunk_size: int = None, **kwargs: Any) -> DataFrame:
-        if chunk_size is not None:
-            return super().df(chunk_size=chunk_size)
-        query_job: bigquery.QueryJob = self.native_cursor._query_job
+    def __init__(self, curr: DBApiCursor) -> None:
+        super().__init__(curr)
+        self.df_iterator = None
 
+    def df(self, chunk_size: int = None, **kwargs: Any) -> DataFrame:
+        query_job: bigquery.QueryJob = self.native_cursor.query_job
+        if self.df_iterator:
+            return next(self.df_iterator, None)
         try:
+            if chunk_size is not None:
+                # create iterator with given page size
+                self.df_iterator = query_job.result(page_size=chunk_size).to_dataframe_iterable()
+                return next(self.df_iterator, None)
             return query_job.to_dataframe(**kwargs)
-        except ValueError:
+        except ValueError as ex:
             # no pyarrow/db-types, fallback to our implementation
-            return super().df()
+            logger.warning(f"Native BigQuery pandas reader could not be used: {str(ex)}")
+            return super().df(chunk_size=chunk_size)
+
+    def close(self) -> None:
+        if self.df_iterator:
+            self.df_iterator.close()
 
 
 class BigQuerySqlClient(SqlClientBase[bigquery.Client], DBTransaction):
@@ -220,12 +234,11 @@ def execute_query(self, query: AnyStr, *args: Any, **kwargs: Any) -> Iterator[DB
                 conn.close()
 
     def fully_qualified_dataset_name(self, escape: bool = True) -> str:
+        project_id = self.capabilities.case_identifier(self.credentials.project_id)
+        dataset_name = self.capabilities.case_identifier(self.dataset_name)
         if escape:
-            project_id = self.capabilities.escape_identifier(self.credentials.project_id)
-            dataset_name = self.capabilities.escape_identifier(self.dataset_name)
-        else:
-            project_id = self.credentials.project_id
-            dataset_name = self.dataset_name
+            project_id = self.capabilities.escape_identifier(project_id)
+            dataset_name = self.capabilities.escape_identifier(dataset_name)
         return f"{project_id}.{dataset_name}"
 
     @classmethod

dlt/destinations/impl/databricks/sql_client.py

@@ -8,8 +8,6 @@
 )
 from databricks.sql.exc import Error as DatabricksSqlError
 
-from dlt.common import pendulum
-from dlt.common import logger
 from dlt.common.destination import DestinationCapabilitiesContext
 from dlt.destinations.exceptions import (
     DatabaseTerminalException,
@@ -22,10 +20,26 @@
     raise_database_error,
     raise_open_connection_error,
 )
-from dlt.destinations.typing import DBApi, DBApiCursor, DBTransaction
+from dlt.destinations.typing import DBApi, DBApiCursor, DBTransaction, DataFrame
 from dlt.destinations.impl.databricks.configuration import DatabricksCredentials
 from dlt.destinations.impl.databricks import capabilities
-from dlt.common.time import to_py_date, to_py_datetime
+
+
+class DatabricksCursorImpl(DBApiCursorImpl):
+    """Use native data frame support if available"""
+
+    native_cursor: DatabricksSqlCursor
+    vector_size: ClassVar[int] = 2048
+
+    def df(self, chunk_size: int = None, **kwargs: Any) -> DataFrame:
+        if chunk_size is None:
+            return self.native_cursor.fetchall_arrow().to_pandas()
+        else:
+            df = self.native_cursor.fetchmany_arrow(chunk_size).to_pandas()
+            if df.shape[0] == 0:
+                return None
+            else:
+                return df
 
 
 class DatabricksSqlClient(SqlClientBase[DatabricksSqlConnection], DBTransaction):
@@ -39,7 +53,9 @@ def __init__(self, dataset_name: str, credentials: DatabricksCredentials) -> Non
 
     def open_connection(self) -> DatabricksSqlConnection:
         conn_params = self.credentials.to_connector_params()
-        self._conn = databricks_lib.connect(**conn_params, schema=self.dataset_name)
+        self._conn = databricks_lib.connect(
+            **conn_params, schema=self.dataset_name, use_inline_params="silent"
+        )
         return self._conn
 
     @raise_open_connection_error
@@ -91,6 +107,7 @@ def execute_sql(
     def execute_query(self, query: AnyStr, *args: Any, **kwargs: Any) -> Iterator[DBApiCursor]:
         curr: DBApiCursor = None
         # TODO: databricks connector 3.0.0 will use :named paramstyle only
+        # NOTE: we were able to use the old style until they get deprecated
         # if args:
         #     keys = [f"arg{i}" for i in range(len(args))]
         #     # Replace position arguments (%s) with named arguments (:arg0, :arg1, ...)
@@ -114,15 +131,14 @@ def execute_query(self, query: AnyStr, *args: Any, **kwargs: Any) -> Iterator[DB
             db_args = None
         with self._conn.cursor() as curr:
             curr.execute(query, db_args)
-            yield DBApiCursorImpl(curr)  # type: ignore[abstract]
+            yield DatabricksCursorImpl(curr)  # type: ignore[abstract]
 
     def fully_qualified_dataset_name(self, escape: bool = True) -> str:
+        catalog = self.capabilities.case_identifier(self.credentials.catalog)
+        dataset_name = self.capabilities.case_identifier(self.dataset_name)
         if escape:
-            catalog = self.capabilities.escape_identifier(self.credentials.catalog)
-            dataset_name = self.capabilities.escape_identifier(self.dataset_name)
-        else:
-            catalog = self.credentials.catalog
-            dataset_name = self.dataset_name
+            catalog = self.capabilities.escape_identifier(catalog)
+            dataset_name = self.capabilities.escape_identifier(dataset_name)
         return f"{catalog}.{dataset_name}"
 
     @staticmethod