Implement transfer for GCS/S3 to BigQuery

universal_transfer_operator/example_dags/example_universal_transfer_operator.py

@@ -3,7 +3,7 @@
 
 from airflow import DAG
 
-from universal_transfer_operator.constants import TransferMode
+from universal_transfer_operator.constants import FileType, TransferMode
 from universal_transfer_operator.datasets.file.base import File
 from universal_transfer_operator.datasets.table import Metadata, Table
 from universal_transfer_operator.integrations.fivetran import Connector, Destination, FiveTranOptions, Group
@@ -30,6 +30,75 @@
         ),
     )
 
+    transfer_non_native_s3_to_sqlite = UniversalTransferOperator(
+        task_id="transfer_non_native_s3_to_sqlite",
+        source_dataset=File(
+            path="s3://astro-sdk-test/uto/csv_files/", conn_id="aws_default", filetype=FileType.CSV
+        ),
+        destination_dataset=Table(name="uto_s3_table", conn_id="sqlite_default"),
+    )
+
+    transfer_non_native_gs_to_sqlite = UniversalTransferOperator(
+        task_id="transfer_non_native_gs_to_sqlite",
+        source_dataset=File(
+            path="gs://uto-test/uto/csv_files/", conn_id="google_cloud_default", filetype=FileType.CSV
+        ),
+        destination_dataset=Table(name="uto_gs_table", conn_id="sqlite_default"),
+    )
+
+    transfer_non_native_s3_to_snowflake = UniversalTransferOperator(
+        task_id="transfer_non_native_s3_to_snowflake",
+        source_dataset=File(
+            path="s3://astro-sdk-test/uto/csv_files/", conn_id="aws_default", filetype=FileType.CSV
+        ),
+        destination_dataset=Table(name="uto_s3_table", conn_id="snowflake_default"),
+    )
+
+    transfer_non_native_gs_to_snowflake = UniversalTransferOperator(
+        task_id="transfer_non_native_gs_to_snowflake",
+        source_dataset=File(
+            path="gs://uto-test/uto/csv_files/", conn_id="google_cloud_default", filetype=FileType.CSV
+        ),
+        destination_dataset=Table(name="uto_gs_table", conn_id="snowflake_default"),
+    )
+
+    transfer_non_native_gs_to_bigquery = UniversalTransferOperator(
+        task_id="transfer_non_native_gs_to_bigquery",
+        source_dataset=File(path="gs://uto-test/uto/homes_main.csv", conn_id="google_cloud_default"),
+        destination_dataset=Table(
+            name="uto_gs_to_bigquery_table", conn_id="google_cloud_default", metadata=Metadata(schema="astro")
+        ),
+    )
+
+    transfer_non_native_s3_to_bigquery = UniversalTransferOperator(
+        task_id="transfer_non_native_s3_to_bigquery",
+        source_dataset=File(
+            path="s3://astro-sdk-test/uto/csv_files/", conn_id="aws_default", filetype=FileType.CSV
+        ),
+        destination_dataset=Table(
+            name="uto_s3_to_bigquery_table", conn_id="google_cloud_default", metadata=Metadata(schema="astro")
+        ),
+    )
+
+    transfer_non_native_bigquery_to_snowflake = UniversalTransferOperator(
+        task_id="transfer_non_native_bigquery_to_snowflake",
+        source_dataset=Table(
+            name="uto_s3_to_bigquery_table", conn_id="google_cloud_default", metadata=Metadata(schema="astro")
+        ),
+        destination_dataset=Table(
+            name="uto_bigquery_to_snowflake_table",
+            conn_id="snowflake_default",
+        ),
+    )
+
+    transfer_non_native_bigquery_to_sqlite = UniversalTransferOperator(
+        task_id="transfer_non_native_bigquery_to_sqlite",
+        source_dataset=Table(
+            name="uto_s3_to_bigquery_table", conn_id="google_cloud_default", metadata=Metadata(schema="astro")
+        ),
+        destination_dataset=Table(name="uto_bigquery_to_sqlite_table", conn_id="sqlite_default"),
+    )
+
     transfer_fivetran_with_connector_id = UniversalTransferOperator(
         task_id="transfer_fivetran_with_connector_id",
         source_dataset=File(path="s3://astro-sdk-test/uto/", conn_id="aws_default"),

universal_transfer_operator/pyproject.toml

@@ -54,7 +54,10 @@ google = [
 snowflake = [
     "apache-airflow-providers-snowflake",
     "snowflake-sqlalchemy>=1.2.0",
-    "snowflake-connector-python[pandas]",
+    "snowflake-connector-python[pandas]<3.0.0",
+    # pinning snowflake-connector-python[pandas]<3.0.0 due to a conflict in snowflake-connector-python/pyarrow/google
+    # packages and pandas-gbq/google packages which is forcing pandas-gbq of version 0.13.2 installed, which is not
+    # compatible with pandas 1.5.3
 ]
 
 amazon = [
@@ -72,7 +75,10 @@ all = [
     "apache-airflow-providers-google>=6.4.0",
     "apache-airflow-providers-snowflake",
     "smart-open[all]>=5.2.1",
-    "snowflake-connector-python[pandas]",
+    "snowflake-connector-python[pandas]<3.0.0",
+    # pinning snowflake-connector-python[pandas]<3.0.0 due to a conflict in snowflake-connector-python/pyarrow/google
+    # packages and pandas-gbq/google packages which is forcing pandas-gbq of version 0.13.2 installed, which is not
+    # compatible with pandas 1.5.3
     "snowflake-sqlalchemy>=1.2.0",
     "sqlalchemy-bigquery>=1.3.0",
     "s3fs",

universal_transfer_operator/src/universal_transfer_operator/constants.py

@@ -99,3 +99,4 @@ def __repr__(self):
 LoadExistStrategy = Literal["replace", "append"]
 DEFAULT_CHUNK_SIZE = 1000000
 ColumnCapitalization = Literal["upper", "lower", "original"]
+DEFAULT_SCHEMA = "tmp_transfers"

universal_transfer_operator/src/universal_transfer_operator/data_providers/__init__.py

@@ -5,12 +5,40 @@
 from universal_transfer_operator.constants import TransferMode
 from universal_transfer_operator.data_providers.base import DataProviders
 from universal_transfer_operator.datasets.base import Dataset
+from universal_transfer_operator.datasets.file.base import File
+from universal_transfer_operator.datasets.table import Table
 from universal_transfer_operator.utils import TransferParameters, get_class_name
 
-DATASET_CONN_ID_TO_DATAPROVIDER_MAPPING = dict.fromkeys(
-    ["s3", "aws"], "universal_transfer_operator.data_providers.filesystem.aws.s3"
-) | dict.fromkeys(
-    ["gs", "google_cloud_platform"], "universal_transfer_operator.data_providers.filesystem.google.cloud.gcs"
+DATASET_CONN_ID_TO_DATAPROVIDER_MAPPING = (
+    dict.fromkeys(
+        [("s3", File), ("aws", File)], "universal_transfer_operator.data_providers.filesystem.aws.s3"
+    )
+    | dict.fromkeys(
+        [("gs", Table), ("google_cloud_platform", Table)],
+        "universal_transfer_operator.data_providers.database.google.bigquery",
+    )
+    | dict.fromkeys(
+        [("gs", File), ("google_cloud_platform", File)],
+        "universal_transfer_operator.data_providers.filesystem.google.cloud.gcs",
+    )
+    | dict.fromkeys(
+        [
+            ("sqlite", Table),
+        ],
+        "universal_transfer_operator.data_providers.database.sqlite",
+    )
+    | dict.fromkeys(
+        [
+            ("snowflake", Table),
+        ],
+        "universal_transfer_operator.data_providers.database.snowflake",
+    )
+    | dict.fromkeys(
+        [
+            (None, File),
+        ],
+        "universal_transfer_operator.data_providers.filesystem.local",
+    )
 )
 
 
@@ -19,8 +47,12 @@ def create_dataprovider(
     transfer_params: TransferParameters = None,
     transfer_mode: TransferMode = TransferMode.NONNATIVE,
 ) -> DataProviders:
-    conn_type = BaseHook.get_connection(dataset.conn_id).conn_type
-    module_path = DATASET_CONN_ID_TO_DATAPROVIDER_MAPPING[conn_type]
+    print(dataset)
+    if dataset.conn_id != "":
+        conn_type = BaseHook.get_connection(dataset.conn_id).conn_type
+    else:
+        conn_type = None
+    module_path = DATASET_CONN_ID_TO_DATAPROVIDER_MAPPING[(conn_type, type(dataset))]
     module = importlib.import_module(module_path)
     class_name = get_class_name(module_ref=module, suffix="DataProvider")
     data_provider: DataProviders = getattr(module, class_name)(

universal_transfer_operator/src/universal_transfer_operator/data_providers/base.py

@@ -97,3 +97,11 @@ def openlineage_dataset_name(self) -> str:
         https://github.com/OpenLineage/OpenLineage/blob/main/spec/Naming.md
         """
         raise NotImplementedError
+
+    @property
+    def openlineage_dataset_uri(self) -> str:
+        """
+        Returns the open lineage dataset uri as per
+        https://github.com/OpenLineage/OpenLineage/blob/main/spec/Naming.md
+        """
+        raise NotImplementedError