Repository: apache/airflow
File: providers/google/src/airflow/providers/google/cloud/transfers/mssql_to_gcs.py (Python)

Code before:
    **Example**:
        The following operator will export data from the Customers table
        within the given MSSQL Database and then upload it to the
        'mssql-export' GCS bucket (along with a schema file). ::

            export_customers = MSSQLToGCSOperator(
                task_id="export_customers",
                sql="SELECT * FROM dbo.Customers;",
                bit_fields=["some_bit_field", "another_bit_field"],
                bucket="mssql-export",
                filename="data/customers/export.json",
                schema_filename="schemas/export.json",
                mssql_conn_id="mssql_default",
                gcp_conn_id="google_cloud_default",
                dag=dag,
            )

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:MSSQLToGCSOperator`

    """

    ui_color = "#e0a98c"

    type_map = {2: "BOOLEAN", 3: "INTEGER", 4: "TIMESTAMP", 5: "NUMERIC"}

    def __init__(
        self,
        *,
        bit_fields: Sequence[str] | None = None,
        mssql_conn_id="mssql_default",
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.mssql_conn_id = mssql_conn_id
        self.bit_fields = bit_fields or []

    @cached_property
    def db_hook(self) -> MsSqlHook:
        return MsSqlHook(mssql_conn_id=self.mssql_conn_id)

    def query(self):
        """
        Query MSSQL and returns a cursor of results.

        :return: mssql cursor
        """
        conn = self.db_hook.get_conn()
        cursor = conn.cursor()
        cursor.execute(self.sql)

Review comment (bug):
@juarezr as I understood this `type_map` variable for JSON, CSV and Parquet data types and if you change this `BOOLEAN` to `BOOL` for Parquet format it means for me that it will fail for JSON and CSV formats. Correct me if I am wrong?

Code after:
    **Example**:
        The following operator will export data from the Customers table
        within the given MSSQL Database and then upload it to the
        'mssql-export' GCS bucket (along with a schema file). ::

            export_customers = MSSQLToGCSOperator(
                task_id="export_customers",
                sql="SELECT * FROM dbo.Customers;",
                bit_fields=["some_bit_field", "another_bit_field"],
                bucket="mssql-export",
                filename="data/customers/export.json",
                schema_filename="schemas/export.json",
                mssql_conn_id="mssql_default",
                gcp_conn_id="google_cloud_default",
                dag=dag,
            )

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:MSSQLToGCSOperator`

    """

    ui_color = "#e0a98c"

    type_map = {2: "BOOL", 3: "INTEGER", 4: "TIMESTAMP", 5: "NUMERIC"}

    def __init__(
        self,
        *,
        bit_fields: Sequence[str] | None = None,
        mssql_conn_id="mssql_default",
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.mssql_conn_id = mssql_conn_id
        self.bit_fields = bit_fields or []

    @cached_property
    def db_hook(self) -> MsSqlHook:
        return MsSqlHook(mssql_conn_id=self.mssql_conn_id)

    def query(self):
        """
        Query MSSQL and returns a cursor of results.

        :return: mssql cursor
        """
        conn = self.db_hook.get_conn()
        cursor = conn.cursor()
        cursor.execute(self.sql)