diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index 57aa4efa466ea..bc7d58d0990b7 100644 --- a/.github/workflows/integration-tests.yml +++ b/.github/workflows/integration-tests.yml @@ -140,6 +140,11 @@ jobs: breeze testing integration-tests --integration qdrant breeze down if: inputs.is-airflow-runner != 'true' && matrix.backend == 'postgres' + - name: "Integration Tests Postgres: mssql" + run: | + breeze testing integration-tests --integration mssql + breeze down + if: inputs.is-airflow-runner != 'true' && matrix.backend == 'postgres' # For Self Hosted runners we run complete set of integration tests # in one big test - as we have enough resources to start all integrations # and we do that for both - postgres and mysql diff --git a/airflow/providers/microsoft/mssql/hooks/mssql.py b/airflow/providers/microsoft/mssql/hooks/mssql.py index a7a35237042cb..a231590ae2ce7 100644 --- a/airflow/providers/microsoft/mssql/hooks/mssql.py +++ b/airflow/providers/microsoft/mssql/hooks/mssql.py @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any import pymssql +from pymssql import Connection as PymssqlConnection from airflow.providers.common.sql.hooks.sql import DbApiHook @@ -104,7 +105,7 @@ def get_sqlalchemy_connection( engine = self.get_sqlalchemy_engine(engine_kwargs=engine_kwargs) return engine.connect(**(connect_kwargs or {})) - def get_conn(self) -> pymssql.connect: + def get_conn(self) -> PymssqlConnection: """Return ``pymssql`` connection object.""" conn = self.connection return pymssql.connect( @@ -112,15 +113,15 @@ def get_conn(self) -> pymssql.connect: user=conn.login, password=conn.password, database=self.schema or conn.schema, - port=conn.port, + port=str(conn.port), ) def set_autocommit( self, - conn: pymssql.connect, + conn: PymssqlConnection, autocommit: bool, ) -> None: conn.autocommit(autocommit) - def get_autocommit(self, conn: pymssql.connect): + def get_autocommit(self, conn: PymssqlConnection): return conn.autocommit_state diff --git a/contributing-docs/testing/integration_tests.rst b/contributing-docs/testing/integration_tests.rst index 8650df1b508cd..0d8a94af99ff7 100644 --- a/contributing-docs/testing/integration_tests.rst +++ b/contributing-docs/testing/integration_tests.rst @@ -58,6 +58,8 @@ The following integrations are available: +--------------+----------------------------------------------------+ | mongo | Integration required for MongoDB hooks. | +--------------+----------------------------------------------------+ +| mssql | Integration required for mssql hooks. | ++--------------+----------------------------------------------------+ | openlineage | Integration required for Openlineage hooks. | +--------------+----------------------------------------------------+ | otel | Integration required for OTEL/opentelemetry hooks. | diff --git a/dev/breeze/doc/images/output-commands.svg b/dev/breeze/doc/images/output-commands.svg index 7ee81cf64ad7b..b8407c5808d1c 100644 --- a/dev/breeze/doc/images/output-commands.svg +++ b/dev/breeze/doc/images/output-commands.svg @@ -299,8 +299,8 @@ │(>3.8< | 3.9 | 3.10 | 3.11 | 3.12)                          │ │[default: 3.8]                                              │ │--integrationIntegration(s) to enable when running (can be more than one).                       │ -│(all | all-testable | cassandra | celery | kafka | kerberos | mongo | openlineage | │ -│otel | pinot | qdrant | statsd | trino)                                             │ +│(all | all-testable | cassandra | celery | kafka | kerberos | mongo | mssql |       │ +│openlineage | otel | pinot | qdrant | statsd | trino)                               │ │--standalone-dag-processorRun standalone dag processor for start-airflow.│ │--database-isolationRun airflow in database isolation mode.│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ diff --git a/dev/breeze/doc/images/output_shell.svg b/dev/breeze/doc/images/output_shell.svg index 0b30f0a04f830..22db8cc492a35 100644 --- a/dev/breeze/doc/images/output_shell.svg +++ b/dev/breeze/doc/images/output_shell.svg @@ -477,138 +477,138 @@ Enter breeze environment. this is the default command use when no other is selected. ╭─ Execution mode ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--python-pPython major/minor version used in Airflow image for images.│ +│--python-pPython major/minor version used in Airflow image for images.│ │(>3.8< | 3.9 | 3.10 | 3.11 | 3.12)                          │ │[default: 3.8]                                              │ -│--integrationIntegration(s) to enable when running (can be more than one).                       │ -│(all | all-testable | cassandra | celery | kafka | kerberos | mongo | openlineage | │ -│otel | pinot | qdrant | statsd | trino)                                             │ -│--standalone-dag-processorRun standalone dag processor for start-airflow.│ -│--database-isolationRun airflow in database isolation mode.│ +│--integrationIntegration(s) to enable when running (can be more than one).                       │ +│(all | all-testable | cassandra | celery | kafka | kerberos | mongo | mssql |       │ +│openlineage | otel | pinot | qdrant | statsd | trino)                               │ +│--standalone-dag-processorRun standalone dag processor for start-airflow.│ +│--database-isolationRun airflow in database isolation mode.│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Docker Compose project management ──────────────────────────────────────────────────────────────────────────────────╮ -│--project-nameName of the docker-compose project to bring down. The `docker-compose` is for legacy  │ -│breeze project name and you can use `breeze down --project-name docker-compose` to    │ +│--project-nameName of the docker-compose project to bring down. The `docker-compose` is for legacy  │ +│breeze project name and you can use `breeze down --project-name docker-compose` to    │ │stop all containers belonging to it.                                                  │ │(breeze | pre-commit | docker-compose)                                                │ │[default: breeze]                                                                     │ -│--restart,--remove-orphansRestart all containers before entering shell (also removes orphan containers).│ -│--docker-hostOptional - docker host to use when running docker commands. When set, the `--builder` │ +│--restart,--remove-orphansRestart all containers before entering shell (also removes orphan containers).│ +│--docker-hostOptional - docker host to use when running docker commands. When set, the `--builder` │ │option is ignored when building images.                                               │ │(TEXT)                                                                                │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Scripts execution ──────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--quietSuppress initialization output when starting.│ -│--skip-image-upgrade-checkSkip checking if the CI image is up to date.│ -│--warn-image-upgrade-neededWarn when image upgrade is needed even if --skip-upgrade-check is used.│ -│--skip-environment-initializationSkip running breeze entrypoint initialization - no user output, no db checks.│ -│--ttyWhether to allocate pseudo-tty when running docker command (useful for         │ +│--quietSuppress initialization output when starting.│ +│--skip-image-upgrade-checkSkip checking if the CI image is up to date.│ +│--warn-image-upgrade-neededWarn when image upgrade is needed even if --skip-upgrade-check is used.│ +│--skip-environment-initializationSkip running breeze entrypoint initialization - no user output, no db checks.│ +│--ttyWhether to allocate pseudo-tty when running docker command (useful for         │ │pre-commit and CI to force-enable it).                                         │ │(auto | enabled | disabled)                                                    │ │[default: auto]                                                                │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Database ───────────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--backend-bDatabase backend to use. If 'none' is selected, breeze starts with invalid DB configuration │ +│--backend-bDatabase backend to use. If 'none' is selected, breeze starts with invalid DB configuration │ │and no database and any attempts to connect to Airflow DB will fail.                        │ │(>sqlite< | mysql | postgres | none)                                                        │ │[default: sqlite]                                                                           │ -│--postgres-version-PVersion of Postgres used.(>12< | 13 | 14 | 15 | 16)[default: 12]│ -│--mysql-version-MVersion of MySQL used.(>8.0< | 8.4)[default: 8.0]│ -│--db-reset-dReset DB when entering the container.│ +│--postgres-version-PVersion of Postgres used.(>12< | 13 | 14 | 15 | 16)[default: 12]│ +│--mysql-version-MVersion of MySQL used.(>8.0< | 8.4)[default: 8.0]│ +│--db-reset-dReset DB when entering the container.│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Choose executor ────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--executorSpecify the executor to use with shell command.                                              │ +│--executorSpecify the executor to use with shell command.                                              │ │(LocalExecutor|KubernetesExecutor|CeleryExecutor|CeleryKubernetesExecutor|SequentialExecutor)│ │[default: LocalExecutor]                                                                     │ -│--celery-brokerSpecify the celery message broker(rabbitmq|redis)[default: redis]│ -│--celery-flowerStart celery flower│ +│--celery-brokerSpecify the celery message broker(rabbitmq|redis)[default: redis]│ +│--celery-flowerStart celery flower│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Build CI image (before entering shell) ─────────────────────────────────────────────────────────────────────────────╮ -│--force-buildForce image build no matter if it is determined as needed.│ -│--platformPlatform for Airflow image.(linux/amd64 | linux/arm64)│ -│--image-tagTag of the image which is used to run the image (implies --mount-sources=skip).(TEXT)│ +│--force-buildForce image build no matter if it is determined as needed.│ +│--platformPlatform for Airflow image.(linux/amd64 | linux/arm64)│ +│--image-tagTag of the image which is used to run the image (implies --mount-sources=skip).(TEXT)│ │[default: latest]                                                              │ -│--github-repository-gGitHub repository used to pull, push run images.(TEXT)[default: apache/airflow]│ -│--builderBuildx builder used to perform `docker buildx build` commands.(TEXT)│ +│--github-repository-gGitHub repository used to pull, push run images.(TEXT)[default: apache/airflow]│ +│--builderBuildx builder used to perform `docker buildx build` commands.(TEXT)│ │[default: autodetect]                                         │ -│--use-uv/--no-use-uvUse uv instead of pip as packaging tool to build the image.[default: use-uv]│ -│--uv-http-timeoutTimeout for requests that UV makes (only used in case of UV builds).(INTEGER RANGE)│ +│--use-uv/--no-use-uvUse uv instead of pip as packaging tool to build the image.[default: use-uv]│ +│--uv-http-timeoutTimeout for requests that UV makes (only used in case of UV builds).(INTEGER RANGE)│ │[default: 300; x>=1]                                                │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Mounting the sources and volumes ───────────────────────────────────────────────────────────────────────────────────╮ -│--mount-sourcesChoose scope of local sources that should be mounted, skipped, or removed (default =       │ +│--mount-sourcesChoose scope of local sources that should be mounted, skipped, or removed (default =       │ │selected).                                                                                 │ │(selected | all | skip | remove | tests)                                                   │ │[default: selected]                                                                        │ -│--include-mypy-volumeWhether to include mounting of the mypy volume (useful for debugging mypy).│ +│--include-mypy-volumeWhether to include mounting of the mypy volume (useful for debugging mypy).│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Installing packages after entering shell ───────────────────────────────────────────────────────────────────────────╮ -│--airflow-constraints-locationLocation of airflow constraints to use (remote URL or   │ +│--airflow-constraints-locationLocation of airflow constraints to use (remote URL or   │ │local context file).                                    │ │(TEXT)                                                  │ -│--airflow-constraints-modeMode of constraints for Airflow for CI image building.  │ +│--airflow-constraints-modeMode of constraints for Airflow for CI image building.  │ │(constraints-source-providers | constraints |           │ │constraints-no-providers)                               │ │[default: constraints-source-providers]                 │ -│--airflow-constraints-referenceConstraint reference to use for airflow installation    │ +│--airflow-constraints-referenceConstraint reference to use for airflow installation    │ │(used in calculated constraints URL).                   │ │(TEXT)                                                  │ -│--airflow-extrasAirflow extras to install when --use-airflow-version is │ +│--airflow-extrasAirflow extras to install when --use-airflow-version is │ │used                                                    │ │(TEXT)                                                  │ -│--airflow-skip-constraintsDo not use constraints when installing airflow.│ -│--install-airflow-with-constraints/--no-install-airflow…Install airflow in a separate step, with constraints    │ +│--airflow-skip-constraintsDo not use constraints when installing airflow.│ +│--install-airflow-with-constraints/--no-install-airflow…Install airflow in a separate step, with constraints    │ │determined from package or airflow version.             │ │[default: install-airflow-with-constraints]             │ -│--install-selected-providersComma-separated list of providers selected to be        │ -│installed (implies --use-packages-from-dist).           │ +│--install-selected-providersComma-separated list of providers selected to be        │ +│installed (implies --use-packages-from-dist).           │ │(TEXT)                                                  │ -│--package-formatFormat of packages that should be installed from dist.│ +│--package-formatFormat of packages that should be installed from dist.│ │(wheel | sdist)                                       │ │[default: wheel]                                      │ -│--providers-constraints-locationLocation of providers constraints to use (remote URL or │ +│--providers-constraints-locationLocation of providers constraints to use (remote URL or │ │local context file).                                    │ │(TEXT)                                                  │ -│--providers-constraints-modeMode of constraints for Providers for CI image building.│ +│--providers-constraints-modeMode of constraints for Providers for CI image building.│ │(constraints-source-providers | constraints |           │ │constraints-no-providers)                               │ │[default: constraints-source-providers]                 │ -│--providers-constraints-referenceConstraint reference to use for providers installation  │ +│--providers-constraints-referenceConstraint reference to use for providers installation  │ │(used in calculated constraints URL). Can be 'default'  │ │in which case the default constraints-reference is used.│ │(TEXT)                                                  │ -│--providers-skip-constraintsDo not use constraints when installing providers.│ -│--use-airflow-versionUse (reinstall at entry) Airflow version from PyPI. It  │ +│--providers-skip-constraintsDo not use constraints when installing providers.│ +│--use-airflow-versionUse (reinstall at entry) Airflow version from PyPI. It  │ │can also be version (to install from PyPI), `none`,     │ │`wheel`, or `sdist` to install from `dist` folder, or   │ │VCS URL to install from                                 │ │(https://pip.pypa.io/en/stable/topics/vcs-support/).    │ -│Implies --mount-sources `remove`.                       │ +│Implies --mount-sources `remove`.                       │ │(none | wheel | sdist | <airflow_version>)              │ -│--use-packages-from-distInstall all found packages (--package-format determines │ +│--use-packages-from-distInstall all found packages (--package-format determines │ │type) from 'dist' folder when entering breeze.          │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Upgrading/downgrading/removing selected packages ───────────────────────────────────────────────────────────────────╮ -│--upgrade-botoRemove aiobotocore and upgrade botocore and boto to the latest version.│ -│--downgrade-sqlalchemyDowngrade SQLAlchemy to minimum supported version.│ -│--downgrade-pendulumDowngrade Pendulum to minimum supported version.│ -│--pydanticDetermines which pydantic should be used during tests.(v2 | v1 | none)[default: v2]│ +│--upgrade-botoRemove aiobotocore and upgrade botocore and boto to the latest version.│ +│--downgrade-sqlalchemyDowngrade SQLAlchemy to minimum supported version.│ +│--downgrade-pendulumDowngrade Pendulum to minimum supported version.│ +│--pydanticDetermines which pydantic should be used during tests.(v2 | v1 | none)[default: v2]│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ DB test flags ──────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--run-db-tests-onlyOnly runs tests that require a database│ -│--skip-db-testsSkip tests that require a database│ +│--run-db-tests-onlyOnly runs tests that require a database│ +│--skip-db-testsSkip tests that require a database│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Other options ──────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--forward-credentials-fForward local credentials to container when running.│ -│--max-timeMaximum time that the command should take - if it takes longer, the command will fail.│ +│--forward-credentials-fForward local credentials to container when running.│ +│--max-timeMaximum time that the command should take - if it takes longer, the command will fail.│ │(INTEGER RANGE)                                                                       │ -│--verbose-commandsShow details of commands executed.│ +│--verbose-commandsShow details of commands executed.│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--answer-aForce answer to questions.(y | n | q | yes | no | quit)│ -│--dry-run-DIf dry-run is set, commands are only printed, not executed.│ -│--verbose-vPrint verbose information about performed steps.│ -│--help-hShow this message and exit.│ +│--answer-aForce answer to questions.(y | n | q | yes | no | quit)│ +│--dry-run-DIf dry-run is set, commands are only printed, not executed.│ +│--verbose-vPrint verbose information about performed steps.│ +│--help-hShow this message and exit.│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ diff --git a/dev/breeze/doc/images/output_shell.txt b/dev/breeze/doc/images/output_shell.txt index a5de913439c37..55d927812786b 100644 --- a/dev/breeze/doc/images/output_shell.txt +++ b/dev/breeze/doc/images/output_shell.txt @@ -1 +1 @@ -78a8e3e896df37cb8516c78e3bb67efb +06b133baab26af756b49a756d77b1f11 diff --git a/dev/breeze/doc/images/output_start-airflow.svg b/dev/breeze/doc/images/output_start-airflow.svg index 1f63f2437fc25..4e15edea7fb59 100644 --- a/dev/breeze/doc/images/output_start-airflow.svg +++ b/dev/breeze/doc/images/output_start-airflow.svg @@ -399,8 +399,8 @@ │[default: 3.8]                                              │ │--platformPlatform for Airflow image.(linux/amd64 | linux/arm64)│ │--integrationIntegration(s) to enable when running (can be more than one).                       │ -│(all | all-testable | cassandra | celery | kafka | kerberos | mongo | openlineage | │ -│otel | pinot | qdrant | statsd | trino)                                             │ +│(all | all-testable | cassandra | celery | kafka | kerberos | mongo | mssql |       │ +│openlineage | otel | pinot | qdrant | statsd | trino)                               │ │--standalone-dag-processorRun standalone dag processor for start-airflow.│ │--database-isolationRun airflow in database isolation mode.│ │--load-example-dags-eEnable configuration to load example DAGs when starting Airflow.│ diff --git a/dev/breeze/doc/images/output_start-airflow.txt b/dev/breeze/doc/images/output_start-airflow.txt index c108c4b754c7a..9299177f5cf5d 100644 --- a/dev/breeze/doc/images/output_start-airflow.txt +++ b/dev/breeze/doc/images/output_start-airflow.txt @@ -1 +1 @@ -c56082b081ff31c535dd8e42e545c2ac +881dcef5b50eced507c70ce3bcdfcace diff --git a/dev/breeze/doc/images/output_testing_integration-tests.svg b/dev/breeze/doc/images/output_testing_integration-tests.svg index 15df809d8cebf..6b15f618db443 100644 --- a/dev/breeze/doc/images/output_testing_integration-tests.svg +++ b/dev/breeze/doc/images/output_testing_integration-tests.svg @@ -209,7 +209,7 @@ ╭─ Test environment ───────────────────────────────────────────────────────────────────────────────────────────────────╮ │--integrationIntegration(s) to enable when running (can be more than one).       │ │(all | all-testable | cassandra | celery | kafka | kerberos | mongo │ -│| openlineage | otel | pinot | qdrant | statsd | trino)             │ +│| mssql | openlineage | otel | pinot | qdrant | statsd | trino)     │ │--backend-bDatabase backend to use. If 'none' is selected, breeze starts with  │ │invalid DB configuration and no database and any attempts to connect│ │to Airflow DB will fail.                                            │ diff --git a/dev/breeze/doc/images/output_testing_integration-tests.txt b/dev/breeze/doc/images/output_testing_integration-tests.txt index 227607f308573..c6e3e77f0bb4d 100644 --- a/dev/breeze/doc/images/output_testing_integration-tests.txt +++ b/dev/breeze/doc/images/output_testing_integration-tests.txt @@ -1 +1 @@ -2931809b87863a80ac792f08254a64a0 +a6682a9e42adba04055a0377b17864fb diff --git a/dev/breeze/doc/images/output_testing_tests.svg b/dev/breeze/doc/images/output_testing_tests.svg index 5e1de91ff2be1..0c613f322aec3 100644 --- a/dev/breeze/doc/images/output_testing_tests.svg +++ b/dev/breeze/doc/images/output_testing_tests.svg @@ -416,115 +416,115 @@ more opinionated test suite execution. ╭─ Select test types to run (tests can also be selected by command args individually) ─────────────────────────────────╮ -│--test-typeType of test to run. With Providers, you can specify tests of which providers     │ +│--test-typeType of test to run. With Providers, you can specify tests of which providers     │ │should be run: `Providers[airbyte,http]` or excluded from the full test suite:    │ │`Providers[-amazon,google]`                                                       │ │(All | Default | API | Always | BranchExternalPython | BranchPythonVenv | CLI |   │ │Core | ExternalPython | Operators | Other | PlainAsserts | Providers | PythonVenv │ │| Serialization | WWW | All-Postgres | All-MySQL | All-Quarantined)               │ │[default: Default]                                                                │ -│--parallel-test-typesSpace separated list of test types used for testing in parallel                   │ +│--parallel-test-typesSpace separated list of test types used for testing in parallel                   │ │(API | Always | BranchExternalPython | BranchPythonVenv | CLI | Core |            │ │ExternalPython | Operators | Other | PlainAsserts | Providers | PythonVenv |      │ │Serialization | WWW)                                                              │ │[default: API Always BranchExternalPython BranchPythonVenv CLI Core ExternalPython│ │Operators Other PlainAsserts Providers[-amazon,google] Providers[amazon]          │ │Providers[google] PythonVenv Serialization WWW]                                   │ -│--excluded-parallel-test-typesSpace separated list of test types that will be excluded from parallel tes runs.  │ +│--excluded-parallel-test-typesSpace separated list of test types that will be excluded from parallel tes runs.  │ │(API | Always | BranchExternalPython | BranchPythonVenv | CLI | Core |            │ │ExternalPython | Operators | Other | PlainAsserts | Providers | PythonVenv |      │ │Serialization | WWW)                                                              │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Test options ───────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--test-timeoutTest timeout in seconds. Set the pytest setup, execution and teardown timeouts to this   │ +│--test-timeoutTest timeout in seconds. Set the pytest setup, execution and teardown timeouts to this   │ │value                                                                                    │ │(INTEGER RANGE)                                                                          │ │[default: 60; x>=0]                                                                      │ -│--enable-coverageEnable coverage capturing for tests in the form of XML files│ -│--collect-onlyCollect tests only, do not run them.│ -│--db-reset-dReset DB when entering the container.│ -│--skip-provider-testsSkip provider tests│ +│--enable-coverageEnable coverage capturing for tests in the form of XML files│ +│--collect-onlyCollect tests only, do not run them.│ +│--db-reset-dReset DB when entering the container.│ +│--skip-provider-testsSkip provider tests│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Selectively run DB or non-DB tests ─────────────────────────────────────────────────────────────────────────────────╮ -│--run-db-tests-onlyOnly runs tests that require a database│ -│--skip-db-testsSkip tests that require a database│ +│--run-db-tests-onlyOnly runs tests that require a database│ +│--skip-db-testsSkip tests that require a database│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Test environment ───────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--integrationIntegration(s) to enable when running (can be more than one).       │ +│--integrationIntegration(s) to enable when running (can be more than one).       │ │(all | all-testable | cassandra | celery | kafka | kerberos | mongo │ -│| openlineage | otel | pinot | qdrant | statsd | trino)             │ -│--backend-bDatabase backend to use. If 'none' is selected, breeze starts with  │ +│| mssql | openlineage | otel | pinot | qdrant | statsd | trino)     │ +│--backend-bDatabase backend to use. If 'none' is selected, breeze starts with  │ │invalid DB configuration and no database and any attempts to connect│ │to Airflow DB will fail.                                            │ │(>sqlite< | mysql | postgres | none)                                │ │[default: sqlite]                                                   │ -│--python-pPython major/minor version used in Airflow image for images.│ +│--python-pPython major/minor version used in Airflow image for images.│ │(>3.8< | 3.9 | 3.10 | 3.11 | 3.12)                          │ │[default: 3.8]                                              │ -│--postgres-version-PVersion of Postgres used.(>12< | 13 | 14 | 15 | 16)[default: 12]│ -│--mysql-version-MVersion of MySQL used.(>8.0< | 8.4)[default: 8.0]│ -│--forward-credentials-fForward local credentials to container when running.│ -│--force-sa-warnings/--no-force-sa-warningsEnable `sqlalchemy.exc.MovedIn20Warning` during the tests runs.│ +│--postgres-version-PVersion of Postgres used.(>12< | 13 | 14 | 15 | 16)[default: 12]│ +│--mysql-version-MVersion of MySQL used.(>8.0< | 8.4)[default: 8.0]│ +│--forward-credentials-fForward local credentials to container when running.│ +│--force-sa-warnings/--no-force-sa-warningsEnable `sqlalchemy.exc.MovedIn20Warning` during the tests runs.│ │[default: force-sa-warnings]                                   │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Options for parallel test commands ─────────────────────────────────────────────────────────────────────────────────╮ -│--run-in-parallelRun the operation in parallel on all or selected subset of parameters.│ -│--use-xdistUse xdist plugin for pytest│ -│--parallelismMaximum number of processes to use while running the operation in parallel.│ +│--run-in-parallelRun the operation in parallel on all or selected subset of parameters.│ +│--use-xdistUse xdist plugin for pytest│ +│--parallelismMaximum number of processes to use while running the operation in parallel.│ │(INTEGER RANGE)                                                            │ │[default: 4; 1<=x<=8]                                                      │ -│--skip-cleanupSkip cleanup of temporary files created during parallel run.│ -│--debug-resourcesWhether to show resource information while running in parallel.│ -│--include-success-outputsWhether to include outputs of successful parallel runs (skipped by default).│ +│--skip-cleanupSkip cleanup of temporary files created during parallel run.│ +│--debug-resourcesWhether to show resource information while running in parallel.│ +│--include-success-outputsWhether to include outputs of successful parallel runs (skipped by default).│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Upgrading/downgrading/removing selected packages ───────────────────────────────────────────────────────────────────╮ -│--upgrade-botoRemove aiobotocore and upgrade botocore and boto to the latest version.│ -│--downgrade-sqlalchemyDowngrade SQLAlchemy to minimum supported version.│ -│--downgrade-pendulumDowngrade Pendulum to minimum supported version.│ -│--pydanticDetermines which pydantic should be used during tests.(v2 | v1 | none)[default: v2]│ -│--remove-arm-packagesRemoves arm packages from the image to test if ARM collection works│ +│--upgrade-botoRemove aiobotocore and upgrade botocore and boto to the latest version.│ +│--downgrade-sqlalchemyDowngrade SQLAlchemy to minimum supported version.│ +│--downgrade-pendulumDowngrade Pendulum to minimum supported version.│ +│--pydanticDetermines which pydantic should be used during tests.(v2 | v1 | none)[default: v2]│ +│--remove-arm-packagesRemoves arm packages from the image to test if ARM collection works│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Advanced flag for tests command ────────────────────────────────────────────────────────────────────────────────────╮ -│--airflow-constraints-referenceConstraint reference to use for airflow installation   │ +│--airflow-constraints-referenceConstraint reference to use for airflow installation   │ │(used in calculated constraints URL).                  │ │(TEXT)                                                 │ -│--github-repository-gGitHub repository used to pull, push run images.(TEXT)│ +│--github-repository-gGitHub repository used to pull, push run images.(TEXT)│ │[default: apache/airflow]                       │ -│--image-tagTag of the image which is used to run the image        │ -│(implies --mount-sources=skip).                        │ +│--image-tagTag of the image which is used to run the image        │ +│(implies --mount-sources=skip).                        │ │(TEXT)                                                 │ │[default: latest]                                      │ -│--install-airflow-with-constraints/--no-install-airflo…Install airflow in a separate step, with constraints   │ +│--install-airflow-with-constraints/--no-install-airflo…Install airflow in a separate step, with constraints   │ │determined from package or airflow version.            │ │[default: no-install-airflow-with-constraints]         │ -│--package-formatFormat of packages.(wheel | sdist | both)│ +│--package-formatFormat of packages.(wheel | sdist | both)│ │[default: wheel]   │ -│--providers-constraints-locationLocation of providers constraints to use (remote URL or│ +│--providers-constraints-locationLocation of providers constraints to use (remote URL or│ │local context file).                                   │ │(TEXT)                                                 │ -│--providers-skip-constraintsDo not use constraints when installing providers.│ -│--use-airflow-versionUse (reinstall at entry) Airflow version from PyPI. It │ +│--providers-skip-constraintsDo not use constraints when installing providers.│ +│--use-airflow-versionUse (reinstall at entry) Airflow version from PyPI. It │ │can also be version (to install from PyPI), `none`,    │ │`wheel`, or `sdist` to install from `dist` folder, or  │ │VCS URL to install from                                │ │(https://pip.pypa.io/en/stable/topics/vcs-support/).   │ -│Implies --mount-sources `remove`.                      │ +│Implies --mount-sources `remove`.                      │ │(none | wheel | sdist | <airflow_version>)             │ -│--use-packages-from-distInstall all found packages (--package-format determines│ +│--use-packages-from-distInstall all found packages (--package-format determines│ │type) from 'dist' folder when entering breeze.         │ -│--mount-sourcesChoose scope of local sources that should be mounted,  │ +│--mount-sourcesChoose scope of local sources that should be mounted,  │ │skipped, or removed (default = selected).              │ │(selected | all | skip | remove | tests)               │ │[default: selected]                                    │ -│--skip-docker-compose-downSkips running docker-compose down after tests│ -│--skip-providersComa separated list of providers to skip when running  │ +│--skip-docker-compose-downSkips running docker-compose down after tests│ +│--skip-providersComa separated list of providers to skip when running  │ │tests                                                  │ │(TEXT)                                                 │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--dry-run-DIf dry-run is set, commands are only printed, not executed.│ -│--verbose-vPrint verbose information about performed steps.│ -│--help-hShow this message and exit.│ +│--dry-run-DIf dry-run is set, commands are only printed, not executed.│ +│--verbose-vPrint verbose information about performed steps.│ +│--help-hShow this message and exit.│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ diff --git a/dev/breeze/doc/images/output_testing_tests.txt b/dev/breeze/doc/images/output_testing_tests.txt index 25be0ad1baddb..a2c8ab0a6b621 100644 --- a/dev/breeze/doc/images/output_testing_tests.txt +++ b/dev/breeze/doc/images/output_testing_tests.txt @@ -1 +1 @@ -3c3217a7eceaa77718af4876622e1b0f +e8d3a79d21f9f690bb81fe1a2921a3ca diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py index edbe2e016abc6..2903d0c26e229 100644 --- a/dev/breeze/src/airflow_breeze/global_constants.py +++ b/dev/breeze/src/airflow_breeze/global_constants.py @@ -51,7 +51,17 @@ ALLOWED_BACKENDS = ["sqlite", "mysql", "postgres", "none"] ALLOWED_PROD_BACKENDS = ["mysql", "postgres"] DEFAULT_BACKEND = ALLOWED_BACKENDS[0] -TESTABLE_INTEGRATIONS = ["cassandra", "celery", "kerberos", "mongo", "pinot", "trino", "kafka", "qdrant"] +TESTABLE_INTEGRATIONS = [ + "cassandra", + "celery", + "kerberos", + "mongo", + "pinot", + "trino", + "kafka", + "qdrant", + "mssql", +] OTHER_INTEGRATIONS = ["statsd", "otel", "openlineage"] ALLOWED_DEBIAN_VERSIONS = ["bookworm", "bullseye"] ALL_INTEGRATIONS = sorted( @@ -219,6 +229,7 @@ def get_default_platform_machine() -> str: FLOWER_HOST_PORT = "25555" REDIS_HOST_PORT = "26379" CELERY_BROKER_URLS_MAP = {"rabbitmq": "amqp://guest:guest@rabbitmq:5672", "redis": "redis://redis:6379/0"} +MSSQL_HOST_PORT = "21433" SQLITE_URL = "sqlite:////root/airflow/sqlite/airflow.db" PYTHONDONTWRITEBYTECODE = True diff --git a/dev/breeze/src/airflow_breeze/params/shell_params.py b/dev/breeze/src/airflow_breeze/params/shell_params.py index ccbf45d441cb6..97717db3b7066 100644 --- a/dev/breeze/src/airflow_breeze/params/shell_params.py +++ b/dev/breeze/src/airflow_breeze/params/shell_params.py @@ -43,6 +43,7 @@ MOUNT_REMOVE, MOUNT_SELECTED, MOUNT_TESTS, + MSSQL_HOST_PORT, MYSQL_HOST_PORT, POSTGRES_HOST_PORT, REDIS_HOST_PORT, @@ -523,6 +524,7 @@ def env_variables_for_docker_commands(self) -> dict[str, str]: _set_var(_env, "PYTHON_MAJOR_MINOR_VERSION", self.python) _set_var(_env, "QUIET", self.quiet) _set_var(_env, "REDIS_HOST_PORT", None, REDIS_HOST_PORT) + _set_var(_env, "MSSQL_HOST_PORT", None, MSSQL_HOST_PORT) _set_var(_env, "REGENERATE_MISSING_DOCS", self.regenerate_missing_docs) _set_var(_env, "REMOVE_ARM_PACKAGES", self.remove_arm_packages) _set_var(_env, "RUN_SYSTEM_TESTS", self.run_system_tests) diff --git a/scripts/ci/docker-compose/integration-mssql.yml b/scripts/ci/docker-compose/integration-mssql.yml new file mode 100644 index 0000000000000..4b9f9bf4da54c --- /dev/null +++ b/scripts/ci/docker-compose/integration-mssql.yml @@ -0,0 +1,47 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +--- +version: "3.8" +services: + mssql: + container_name: mssql + image: "mcr.microsoft.com/mssql/server:2019-latest" + labels: + breeze.description: "Integration required for mssql hooks." + volumes: + - mssql-db-volume:/var/opt/mssql + environment: + ACCEPT_EULA: "Y" + MSSQL_PID: "Developer" + MSSQL_SA_PASSWORD: "airflow123" + restart: "on-failure" + ports: + - "${MSSQL_HOST_PORT}:1433" + healthcheck: + test: /opt/mssql-tools/bin/sqlcmd -S localhost,1433 -U sa -P "airflow123" -Q "SELECT 1" -b -o /dev/null + interval: 10s + timeout: 3s + retries: 10 + start_period: 10s + airflow: + depends_on: + mssql: + condition: service_healthy + environment: + - INTEGRATION_MSSQL=true +volumes: + mssql-db-volume: diff --git a/scripts/in_container/check_environment.sh b/scripts/in_container/check_environment.sh index 73686e1368b72..1e9409a502e9f 100755 --- a/scripts/in_container/check_environment.sh +++ b/scripts/in_container/check_environment.sh @@ -188,6 +188,10 @@ if [[ ${INTEGRATION_KAFKA} == "true" ]]; then check_service "Kafka Cluster" "run_nc broker 9092" 50 fi +if [[ ${INTEGRATION_MSSQL} == "true" ]]; then + check_service "mssql" "run_nc mssql 1433" 50 +fi + if [[ ${EXIT_CODE} != 0 ]]; then echo echo "Error: some of the CI environment failed to initialize!" diff --git a/tests/always/test_project_structure.py b/tests/always/test_project_structure.py index f382f0f5ae782..26f609a6f09ea 100644 --- a/tests/always/test_project_structure.py +++ b/tests/always/test_project_structure.py @@ -82,7 +82,6 @@ def test_providers_modules_should_have_tests(self): "tests/providers/apache/hdfs/hooks/test_hdfs.py", "tests/providers/apache/hdfs/log/test_hdfs_task_handler.py", "tests/providers/apache/hdfs/sensors/test_hdfs.py", - "tests/providers/apache/hive/transfers/test_mssql_to_hive.py", "tests/providers/apache/hive/plugins/test_hive.py", "tests/providers/apache/kafka/hooks/test_base.py", "tests/providers/celery/executors/test_celery_executor_utils.py", @@ -143,9 +142,7 @@ def test_providers_modules_should_have_tests(self): "tests/providers/google/cloud/operators/vertex_ai/test_pipeline_job.py", "tests/providers/google/cloud/sensors/test_dataform.py", "tests/providers/google/cloud/transfers/test_bigquery_to_sql.py", - "tests/providers/google/cloud/transfers/test_mssql_to_gcs.py", "tests/providers/google/cloud/transfers/test_presto_to_gcs.py", - "tests/providers/google/cloud/transfers/test_trino_to_gcs.py", "tests/providers/google/cloud/utils/test_bigquery.py", "tests/providers/google/cloud/utils/test_bigquery_get_data.py", "tests/providers/google/cloud/utils/test_dataform.py", diff --git a/tests/integration/providers/apache/hive/transfers/test_mssql_to_hive.py b/tests/integration/providers/apache/hive/transfers/test_mssql_to_hive.py index 876880a1d2eb1..ca62b264ba05a 100644 --- a/tests/integration/providers/apache/hive/transfers/test_mssql_to_hive.py +++ b/tests/integration/providers/apache/hive/transfers/test_mssql_to_hive.py @@ -17,101 +17,68 @@ # under the License. from __future__ import annotations -from unittest.mock import Mock, PropertyMock, patch +import os +from unittest.mock import Mock, PropertyMock, call, patch import pytest -try: - import pymssql +from airflow.providers.apache.hive.transfers.mssql_to_hive import MsSqlToHiveOperator +from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook - from airflow.providers.apache.hive.transfers.mssql_to_hive import MsSqlToHiveOperator -except ImportError: - pytest.skip("MSSQL not available", allow_module_level=True) +AIRFLOW_CONN_MSSQL_DEFAULT = "mssql://sa:airflow123@mssql:1433/" +TEST_TABLE_ID = "Persons" -@pytest.mark.skip(reason="FIXME: Supposed to use 'mssql' integration which is not implemented yet.") @pytest.mark.integration("mssql") class TestMsSqlToHiveTransfer: - def setup_method(self): - self.kwargs = dict(sql="sql", hive_table="table", task_id="test_mssql_to_hive", dag=None) - - def test_type_map_binary(self): - mapped_type = MsSqlToHiveOperator(**self.kwargs).type_map(pymssql.BINARY.value) - - assert mapped_type == "INT" - - def test_type_map_decimal(self): - mapped_type = MsSqlToHiveOperator(**self.kwargs).type_map(pymssql.DECIMAL.value) - - assert mapped_type == "FLOAT" - - def test_type_map_number(self): - mapped_type = MsSqlToHiveOperator(**self.kwargs).type_map(pymssql.NUMBER.value) - - assert mapped_type == "INT" - - def test_type_map_string(self): - mapped_type = MsSqlToHiveOperator(**self.kwargs).type_map(None) + def setup_method(self, mocker): + os.environ["AIRFLOW_CONN_MSSQL_DEFAULT"] = AIRFLOW_CONN_MSSQL_DEFAULT + hook = MsSqlHook() + conn = hook.get_conn() + hook.set_autocommit(conn, True) + self.cursor = conn.cursor() + self.cursor.execute(f"""CREATE TABLE {TEST_TABLE_ID} ( + PersonID int, + LastName varchar(255), + FirstName varchar(255), + Address varchar(255), + City varchar(255) + )""") + self.kwargs = dict( + sql=f"SELECT * FROM {TEST_TABLE_ID}", hive_table="table", task_id="test_mssql_to_hive", dag=None + ) - assert mapped_type == "STRING" + def teardown_method(self): + self.cursor.execute(f"DROP TABLE {TEST_TABLE_ID}") @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.csv") @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.NamedTemporaryFile") - @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.MsSqlHook") @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.HiveCliHook") - def test_execute(self, mock_hive_hook, mock_mssql_hook, mock_tmp_file, mock_csv): + def test_execute(self, mock_hive_hook, mock_tmp_file, mock_csv): type(mock_tmp_file).name = PropertyMock(return_value="tmp_file") mock_tmp_file.return_value.__enter__ = Mock(return_value=mock_tmp_file) - mock_mssql_hook_get_conn = mock_mssql_hook.return_value.get_conn.return_value.__enter__ - mock_mssql_hook_cursor = mock_mssql_hook_get_conn.return_value.cursor.return_value.__enter__ - mock_mssql_hook_cursor.return_value.description = [("anything", "some-other-thing")] - mssql_to_hive_transfer = MsSqlToHiveOperator(**self.kwargs) mssql_to_hive_transfer.execute(context={}) - mock_mssql_hook_cursor.return_value.execute.assert_called_once_with(mssql_to_hive_transfer.sql) mock_tmp_file.assert_called_with(mode="w", encoding="utf-8") mock_csv.writer.assert_called_once_with(mock_tmp_file, delimiter=mssql_to_hive_transfer.delimiter) - field_dict = {} - for field in mock_mssql_hook_cursor.return_value.description: - field_dict[field[0]] = mssql_to_hive_transfer.type_map(field[1]) - mock_csv.writer.return_value.writerows.assert_called_once_with(mock_mssql_hook_cursor.return_value) - mock_hive_hook.return_value.load_file.assert_called_once_with( - mock_tmp_file.name, - mssql_to_hive_transfer.hive_table, - field_dict=field_dict, - create=mssql_to_hive_transfer.create, - partition=mssql_to_hive_transfer.partition, - delimiter=mssql_to_hive_transfer.delimiter, - recreate=mssql_to_hive_transfer.recreate, - tblproperties=mssql_to_hive_transfer.tblproperties, - ) - - @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.csv") - @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.NamedTemporaryFile") - @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.MsSqlHook") - @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.HiveCliHook") - def test_execute_empty_description_field(self, mock_hive_hook, mock_mssql_hook, mock_tmp_file, mock_csv): - type(mock_tmp_file).name = PropertyMock(return_value="tmp_file") - mock_tmp_file.return_value.__enter__ = Mock(return_value=mock_tmp_file) - mock_mssql_hook_get_conn = mock_mssql_hook.return_value.get_conn.return_value.__enter__ - mock_mssql_hook_cursor = mock_mssql_hook_get_conn.return_value.cursor.return_value.__enter__ - mock_mssql_hook_cursor.return_value.description = [("", "")] - - mssql_to_hive_transfer = MsSqlToHiveOperator(**self.kwargs) - mssql_to_hive_transfer.execute(context={}) - - field_dict = {} - for col_count, field in enumerate(mock_mssql_hook_cursor.return_value.description, start=1): - col_position = f"Column{col_count}" - field_dict[col_position] = mssql_to_hive_transfer.type_map(field[1]) - mock_hive_hook.return_value.load_file.assert_called_once_with( - mock_tmp_file.name, - mssql_to_hive_transfer.hive_table, - field_dict=field_dict, - create=mssql_to_hive_transfer.create, - partition=mssql_to_hive_transfer.partition, - delimiter=mssql_to_hive_transfer.delimiter, - recreate=mssql_to_hive_transfer.recreate, - tblproperties=mssql_to_hive_transfer.tblproperties, + mock_hive_hook.return_value.load_file.assert_has_calls( + [ + call( + "tmp_file", + "table", + field_dict={ + "PersonID": "INT", + "LastName": "STRING", + "FirstName": "STRING", + "Address": "STRING", + "City": "STRING", + }, + create=True, + partition={}, + delimiter="\x01", + recreate=False, + tblproperties=None, + ) + ] ) diff --git a/tests/integration/providers/google/cloud/transfers/test_bigquery_to_mssql.py b/tests/integration/providers/google/cloud/transfers/test_bigquery_to_mssql.py index 19ae3dd8a6a21..04271cca69ce6 100644 --- a/tests/integration/providers/google/cloud/transfers/test_bigquery_to_mssql.py +++ b/tests/integration/providers/google/cloud/transfers/test_bigquery_to_mssql.py @@ -17,41 +17,76 @@ # under the License. from __future__ import annotations +import os from unittest import mock import pytest +from google.cloud.bigquery import Row -try: - from airflow.providers.google.cloud.transfers.bigquery_to_mssql import BigQueryToMsSqlOperator -except ImportError: - pytest.skip("MSSQL not available", allow_module_level=True) - +from airflow.providers.google.cloud.transfers.bigquery_to_mssql import BigQueryToMsSqlOperator +from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook +AIRFLOW_CONN_MSSQL_DEFAULT = "mssql://sa:airflow123@mssql:1433/" TASK_ID = "test-bq-create-table-operator" TEST_PROJECT_ID = "test-project" -TEST_DATASET = "test-dataset" -TEST_TABLE_ID = "test-table-id" -TEST_DAG_ID = "test-bigquery-operators" +TEST_DATASET = "test_dataset" +TEST_TABLE_ID = "Persons" +TEST_ROWS = [ + Row((0, "Airflow", "Apache", "1000 N West Street, Suite 1200", "Wilmington, NC, USA"), field_to_index={}) +] -@pytest.mark.skip(reason="FIXME: Supposed to use 'mssql' integration which is not implemented yet.") @pytest.mark.integration("mssql") class TestBigQueryToMsSqlOperator: - @mock.patch("airflow.providers.google.cloud.transfers.bigquery_to_sql.BigQueryHook") - def test_execute_good_request_to_bq(self, mock_hook): - destination_table = "table" + def setup_method(self): + os.environ["AIRFLOW_CONN_MSSQL_DEFAULT"] = AIRFLOW_CONN_MSSQL_DEFAULT + mssql_hook = MsSqlHook() + mssql_conn = mssql_hook.get_conn() + mssql_hook.set_autocommit(mssql_conn, True) + self.mssql_cursor = mssql_conn.cursor() + self.mssql_cursor.execute(f"""CREATE TABLE {TEST_TABLE_ID} ( + PersonID int, + LastName varchar(255), + FirstName varchar(255), + Address varchar(255), + City varchar(255) + )""") + + def teardown_method(self): + self.mssql_cursor.execute(f"DROP TABLE {TEST_TABLE_ID}") + + @mock.patch( + "airflow.providers.google.cloud.transfers.bigquery_to_sql.BigQueryHook", + return_value=mock.MagicMock(list_rows=mock.MagicMock(side_effect=[TEST_ROWS, []])), + ) + def test_execute(self, mock_hook): operator = BigQueryToMsSqlOperator( task_id=TASK_ID, source_project_dataset_table=f"{TEST_PROJECT_ID}.{TEST_DATASET}.{TEST_TABLE_ID}", - target_table_name=destination_table, + target_table_name=TEST_TABLE_ID, + selected_fields=["PersonID", "LastName", "FirstName", "Address", "City"], replace=False, ) - operator.execute(context=mock.MagicMock()) - mock_hook.return_value.list_rows.assert_called_once_with( - dataset_id=TEST_DATASET, - table_id=TEST_TABLE_ID, - max_results=1000, - selected_fields=None, - start_index=0, + mock_hook.return_value.list_rows.assert_has_calls( + [ + mock.call( + dataset_id=TEST_DATASET, + table_id=TEST_TABLE_ID, + max_results=1000, + selected_fields=["PersonID", "LastName", "FirstName", "Address", "City"], + start_index=0, + ), + mock.call( + dataset_id=TEST_DATASET, + table_id=TEST_TABLE_ID, + max_results=1000, + selected_fields=["PersonID", "LastName", "FirstName", "Address", "City"], + start_index=1000, + ), + ], + any_order=False, ) + self.mssql_cursor.execute(f"SELECT * FROM {TEST_TABLE_ID}") + res = self.mssql_cursor.fetchone() + assert res == TEST_ROWS[0].values() diff --git a/tests/integration/providers/google/cloud/transfers/test_mssql_to_gcs.py b/tests/integration/providers/google/cloud/transfers/test_mssql_to_gcs.py index 586cce433847a..919d5387a2afd 100644 --- a/tests/integration/providers/google/cloud/transfers/test_mssql_to_gcs.py +++ b/tests/integration/providers/google/cloud/transfers/test_mssql_to_gcs.py @@ -17,102 +17,65 @@ # under the License. from __future__ import annotations -import datetime +import os from unittest import mock import pytest -try: - from airflow.providers.google.cloud.transfers.mssql_to_gcs import MSSQLToGCSOperator -except ImportError: - pytest.skip("MSSQL not available", allow_module_level=True) +from airflow.providers.google.cloud.transfers.mssql_to_gcs import MSSQLToGCSOperator +from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook +AIRFLOW_CONN_MSSQL_DEFAULT = "mssql://sa:airflow123@mssql:1433/" TASK_ID = "test-mssql-to-gcs" -MSSQL_CONN_ID = "mssql_conn_test" -SQL = "select 1" +TEST_TABLE_ID = "test_table_id" +MSSQL_CONN_ID = "mssql_default" BUCKET = "gs://test" JSON_FILENAME = "test_{}.ndjson" GZIP = False -ROWS = [ - ("mock_row_content_1", 42, True, True), - ("mock_row_content_2", 43, False, False), - ("mock_row_content_3", 44, True, True), -] -CURSOR_DESCRIPTION = ( - ("some_str", 0, None, None, None, None, None), - ("some_num", 3, None, None, None, None, None), - ("some_binary", 2, None, None, None, None, None), - ("some_bit", 3, None, None, None, None, None), -) NDJSON_LINES = [ - b'{"some_binary": true, "some_bit": true, "some_num": 42, "some_str": "mock_row_content_1"}\n', - b'{"some_binary": false, "some_bit": false, "some_num": 43, "some_str": "mock_row_content_2"}\n', - b'{"some_binary": true, "some_bit": true, "some_num": 44, "some_str": "mock_row_content_3"}\n', -] -SCHEMA_FILENAME = "schema_test.json" -SCHEMA_JSON = [ - b'[{"mode": "NULLABLE", "name": "some_str", "type": "STRING"}, ', - b'{"mode": "NULLABLE", "name": "some_num", "type": "INTEGER"}, ', - b'{"mode": "NULLABLE", "name": "some_binary", "type": "BOOLEAN"}, ', - b'{"mode": "NULLABLE", "name": "some_bit", "type": "BOOLEAN"}]', -] - -SCHEMA_JSON_BIT_FIELDS = [ - b'[{"mode": "NULLABLE", "name": "some_str", "type": "STRING"}, ', - b'{"mode": "NULLABLE", "name": "some_num", "type": "INTEGER"}, ', - b'{"mode": "NULLABLE", "name": "some_binary", "type": "BOOLEAN"}, ', - b'{"mode": "NULLABLE", "name": "some_bit", "type": "INTEGER"}]', + b'{"Address": "1000 N West Street, Suite 1200", "City": "Wilmington, NC, USA", "FirstName": "Apache",' + b' "LastName": "Airflow", "PersonID": 0}\n', ] -@pytest.mark.skip(reason="FIXME: Supposed to use 'mssql' integration which is not implemented yet.") @pytest.mark.integration("mssql") class TestMsSqlToGoogleCloudStorageOperator: - @pytest.mark.parametrize( - "value, expected", - [ - ("string", "string"), - (32.9, 32.9), - (-2, -2), - (datetime.date(1970, 1, 2), "1970-01-02"), - (datetime.date(1000, 1, 2), "1000-01-02"), - (datetime.datetime(1970, 1, 1, 1, 0), "1970-01-01T01:00:00"), - (datetime.time(hour=0, minute=0, second=0), "00:00:00"), - (datetime.time(hour=23, minute=59, second=59), "23:59:59"), - ], - ) - def test_convert_type(self, value, expected): + def setup_method(self): + os.environ["AIRFLOW_CONN_MSSQL_DEFAULT"] = AIRFLOW_CONN_MSSQL_DEFAULT + hook = MsSqlHook() + conn = hook.get_conn() + hook.set_autocommit(conn, True) + self.cursor = conn.cursor() + self.cursor.execute(f"""CREATE TABLE {TEST_TABLE_ID} ( + PersonID int, + LastName varchar(255), + FirstName varchar(255), + Address varchar(255), + City varchar(255) + )""") + q = f"""INSERT INTO {TEST_TABLE_ID} ( + PersonID, LastName, FirstName, Address, City + ) VALUES + (0, 'Airflow', 'Apache', '1000 N West Street, Suite 1200', 'Wilmington, NC, USA') + """ + # raise Exception(q) + self.cursor.execute(q) + + def teardown_method(self): + self.cursor.execute(f"DROP TABLE {TEST_TABLE_ID}") + + @mock.patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + def test_execute(self, gcs_hook_mock_class): + """Test successful run of execute function for JSON""" op = MSSQLToGCSOperator( task_id=TASK_ID, mssql_conn_id=MSSQL_CONN_ID, - sql=SQL, + sql=f"SELECT * FROM {TEST_TABLE_ID}", bucket=BUCKET, filename=JSON_FILENAME, ) - assert op.convert_type(value, None) == expected - - def test_init(self): - """Test MySqlToGoogleCloudStorageOperator instance is properly initialized.""" - op = MSSQLToGCSOperator(task_id=TASK_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME) - assert op.task_id == TASK_ID - assert op.sql == SQL - assert op.bucket == BUCKET - assert op.filename == JSON_FILENAME - - @mock.patch("airflow.providers.google.cloud.transfers.mssql_to_gcs.MsSqlHook") - @mock.patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") - def test_exec_success_json(self, gcs_hook_mock_class, mssql_hook_mock_class): - """Test successful run of execute function for JSON""" - op = MSSQLToGCSOperator( - task_id=TASK_ID, mssql_conn_id=MSSQL_CONN_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME - ) - - mssql_hook_mock = mssql_hook_mock_class.return_value - mssql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) - mssql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION - gcs_hook_mock = gcs_hook_mock_class.return_value def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False, metadata=None): @@ -124,73 +87,4 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False, metada assert b"".join(NDJSON_LINES) == file.read() gcs_hook_mock.upload.side_effect = _assert_upload - - op.execute(None) - - mssql_hook_mock_class.assert_called_once_with(mssql_conn_id=MSSQL_CONN_ID) - mssql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL) - - @mock.patch("airflow.providers.google.cloud.transfers.mssql_to_gcs.MsSqlHook") - @mock.patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") - def test_file_splitting(self, gcs_hook_mock_class, mssql_hook_mock_class): - """Test that ndjson is split by approx_max_file_size_bytes param.""" - mssql_hook_mock = mssql_hook_mock_class.return_value - mssql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) - mssql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION - - gcs_hook_mock = gcs_hook_mock_class.return_value - expected_upload = { - JSON_FILENAME.format(0): b"".join(NDJSON_LINES[:2]), - JSON_FILENAME.format(1): NDJSON_LINES[2], - } - - def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False, metadata=None): - assert BUCKET == bucket - assert "application/json" == mime_type - assert GZIP == gzip - with open(tmp_filename, "rb") as file: - assert expected_upload[obj] == file.read() - - gcs_hook_mock.upload.side_effect = _assert_upload - - op = MSSQLToGCSOperator( - task_id=TASK_ID, - sql=SQL, - bucket=BUCKET, - filename=JSON_FILENAME, - approx_max_file_size_bytes=len(expected_upload[JSON_FILENAME.format(0)]), - ) - op.execute(None) - - @mock.patch("airflow.providers.google.cloud.transfers.mssql_to_gcs.MsSqlHook") - @mock.patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") - @pytest.mark.parametrize( - "bit_fields,schema_json", [(None, SCHEMA_JSON), (["bit_fields", SCHEMA_JSON_BIT_FIELDS])] - ) - def test_schema_file(self, gcs_hook_mock_class, mssql_hook_mock_class, bit_fields, schema_json): - """Test writing schema files.""" - mssql_hook_mock = mssql_hook_mock_class.return_value - mssql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) - mssql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION - - gcs_hook_mock = gcs_hook_mock_class.return_value - - def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip, metadata=None): - if obj == SCHEMA_FILENAME: - with open(tmp_filename, "rb") as file: - assert b"".join(SCHEMA_JSON) == file.read() - - gcs_hook_mock.upload.side_effect = _assert_upload - - op = MSSQLToGCSOperator( - task_id=TASK_ID, - sql=SQL, - bucket=BUCKET, - filename=JSON_FILENAME, - schema_filename=SCHEMA_FILENAME, - bit_fields=["some_bit"], - ) - op.execute(None) - - # once for the file and once for the schema - assert 2 == gcs_hook_mock.upload.call_count + op.execute(context=mock.MagicMock()) diff --git a/tests/integration/providers/microsoft/__init__.py b/tests/integration/providers/microsoft/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/tests/integration/providers/microsoft/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/tests/integration/providers/microsoft/mssql/__init__.py b/tests/integration/providers/microsoft/mssql/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/tests/integration/providers/microsoft/mssql/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/tests/integration/providers/microsoft/mssql/hooks/__init__.py b/tests/integration/providers/microsoft/mssql/hooks/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/tests/integration/providers/microsoft/mssql/hooks/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/tests/integration/providers/microsoft/mssql/hooks/test_mssql.py b/tests/integration/providers/microsoft/mssql/hooks/test_mssql.py new file mode 100644 index 0000000000000..ec7c2759e5107 --- /dev/null +++ b/tests/integration/providers/microsoft/mssql/hooks/test_mssql.py @@ -0,0 +1,43 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import os + +import pytest +from pymssql import Connection + +from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook + +AIRFLOW_CONN_MSSQL_DEFAULT = "mssql://sa:airflow123@mssql:1433/" + + +@pytest.mark.integration("mssql") +class TestMsSqlHook: + def setup_method(self): + os.environ["AIRFLOW_CONN_MSSQL_DEFAULT"] = AIRFLOW_CONN_MSSQL_DEFAULT + self.hook = MsSqlHook() + + def test_get_conn(self): + assert isinstance(self.hook.get_conn(), Connection) + + def test_autocommit(self): + conn = self.hook.get_conn() + self.hook.set_autocommit(conn=conn, autocommit=True) + assert self.hook.get_autocommit(conn=conn) + self.hook.set_autocommit(conn=conn, autocommit=False) + assert not self.hook.get_autocommit(conn=conn) diff --git a/tests/providers/apache/hive/transfers/test_mssql_to_hive.py b/tests/providers/apache/hive/transfers/test_mssql_to_hive.py new file mode 100644 index 0000000000000..1fad9b5f8924e --- /dev/null +++ b/tests/providers/apache/hive/transfers/test_mssql_to_hive.py @@ -0,0 +1,110 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest.mock import Mock, PropertyMock, patch + +import pymssql + +from airflow.providers.apache.hive.transfers.mssql_to_hive import MsSqlToHiveOperator + + +class TestMsSqlToHiveTransfer: + def setup_method(self): + self.kwargs = dict(sql="sql", hive_table="table", task_id="test_mssql_to_hive", dag=None) + + def test_type_map_binary(self): + mapped_type = MsSqlToHiveOperator(**self.kwargs).type_map(pymssql.BINARY.value) + + assert mapped_type == "INT" + + def test_type_map_decimal(self): + mapped_type = MsSqlToHiveOperator(**self.kwargs).type_map(pymssql.DECIMAL.value) + + assert mapped_type == "FLOAT" + + def test_type_map_number(self): + mapped_type = MsSqlToHiveOperator(**self.kwargs).type_map(pymssql.NUMBER.value) + + assert mapped_type == "INT" + + def test_type_map_string(self): + mapped_type = MsSqlToHiveOperator(**self.kwargs).type_map(None) + + assert mapped_type == "STRING" + + @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.csv") + @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.NamedTemporaryFile") + @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.MsSqlHook") + @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.HiveCliHook") + def test_execute(self, mock_hive_hook, mock_mssql_hook, mock_tmp_file, mock_csv): + type(mock_tmp_file).name = PropertyMock(return_value="tmp_file") + mock_tmp_file.return_value.__enter__ = Mock(return_value=mock_tmp_file) + mock_mssql_hook_get_conn = mock_mssql_hook.return_value.get_conn.return_value.__enter__ + mock_mssql_hook_cursor = mock_mssql_hook_get_conn.return_value.cursor.return_value.__enter__ + mock_mssql_hook_cursor.return_value.description = [("anything", "some-other-thing")] + + mssql_to_hive_transfer = MsSqlToHiveOperator(**self.kwargs) + mssql_to_hive_transfer.execute(context={}) + + mock_mssql_hook_cursor.return_value.execute.assert_called_once_with(mssql_to_hive_transfer.sql) + mock_tmp_file.assert_called_with(mode="w", encoding="utf-8") + mock_csv.writer.assert_called_once_with(mock_tmp_file, delimiter=mssql_to_hive_transfer.delimiter) + field_dict = {} + for field in mock_mssql_hook_cursor.return_value.description: + field_dict[field[0]] = mssql_to_hive_transfer.type_map(field[1]) + mock_csv.writer.return_value.writerows.assert_called_once_with(mock_mssql_hook_cursor.return_value) + mock_hive_hook.return_value.load_file.assert_called_once_with( + mock_tmp_file.name, + mssql_to_hive_transfer.hive_table, + field_dict=field_dict, + create=mssql_to_hive_transfer.create, + partition=mssql_to_hive_transfer.partition, + delimiter=mssql_to_hive_transfer.delimiter, + recreate=mssql_to_hive_transfer.recreate, + tblproperties=mssql_to_hive_transfer.tblproperties, + ) + + @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.csv") + @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.NamedTemporaryFile") + @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.MsSqlHook") + @patch("airflow.providers.apache.hive.transfers.mssql_to_hive.HiveCliHook") + def test_execute_empty_description_field(self, mock_hive_hook, mock_mssql_hook, mock_tmp_file, mock_csv): + type(mock_tmp_file).name = PropertyMock(return_value="tmp_file") + mock_tmp_file.return_value.__enter__ = Mock(return_value=mock_tmp_file) + mock_mssql_hook_get_conn = mock_mssql_hook.return_value.get_conn.return_value.__enter__ + mock_mssql_hook_cursor = mock_mssql_hook_get_conn.return_value.cursor.return_value.__enter__ + mock_mssql_hook_cursor.return_value.description = [("", "")] + + mssql_to_hive_transfer = MsSqlToHiveOperator(**self.kwargs) + mssql_to_hive_transfer.execute(context={}) + + field_dict = {} + for col_count, field in enumerate(mock_mssql_hook_cursor.return_value.description, start=1): + col_position = f"Column{col_count}" + field_dict[col_position] = mssql_to_hive_transfer.type_map(field[1]) + mock_hive_hook.return_value.load_file.assert_called_once_with( + mock_tmp_file.name, + mssql_to_hive_transfer.hive_table, + field_dict=field_dict, + create=mssql_to_hive_transfer.create, + partition=mssql_to_hive_transfer.partition, + delimiter=mssql_to_hive_transfer.delimiter, + recreate=mssql_to_hive_transfer.recreate, + tblproperties=mssql_to_hive_transfer.tblproperties, + ) diff --git a/tests/providers/google/cloud/transfers/test_mssql_to_gcs.py b/tests/providers/google/cloud/transfers/test_mssql_to_gcs.py new file mode 100644 index 0000000000000..5ee8efe5a45ae --- /dev/null +++ b/tests/providers/google/cloud/transfers/test_mssql_to_gcs.py @@ -0,0 +1,190 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import datetime +from unittest import mock + +import pytest + +from airflow.providers.google.cloud.transfers.mssql_to_gcs import MSSQLToGCSOperator + +TASK_ID = "test-mssql-to-gcs" +MSSQL_CONN_ID = "mssql_conn_test" +SQL = "select 1" +BUCKET = "gs://test" +JSON_FILENAME = "test_{}.ndjson" +GZIP = False + +ROWS = [ + ("mock_row_content_1", 42, True, True), + ("mock_row_content_2", 43, False, False), + ("mock_row_content_3", 44, True, True), +] +CURSOR_DESCRIPTION = ( + ("some_str", 0, None, None, None, None, None), + ("some_num", 3, None, None, None, None, None), + ("some_binary", 2, None, None, None, None, None), + ("some_bit", 3, None, None, None, None, None), +) +NDJSON_LINES = [ + b'{"some_binary": true, "some_bit": true, "some_num": 42, "some_str": "mock_row_content_1"}\n', + b'{"some_binary": false, "some_bit": false, "some_num": 43, "some_str": "mock_row_content_2"}\n', + b'{"some_binary": true, "some_bit": true, "some_num": 44, "some_str": "mock_row_content_3"}\n', +] +SCHEMA_FILENAME = "schema_test.json" +SCHEMA_JSON = [ + b'[{"mode": "NULLABLE", "name": "some_str", "type": "STRING"}, ', + b'{"mode": "NULLABLE", "name": "some_num", "type": "INTEGER"}, ', + b'{"mode": "NULLABLE", "name": "some_binary", "type": "BOOLEAN"}, ', + b'{"mode": "NULLABLE", "name": "some_bit", "type": "BOOLEAN"}]', +] + +SCHEMA_JSON_BIT_FIELDS = [ + b'[{"mode": "NULLABLE", "name": "some_str", "type": "STRING"}, ', + b'{"mode": "NULLABLE", "name": "some_num", "type": "INTEGER"}, ', + b'{"mode": "NULLABLE", "name": "some_binary", "type": "BOOLEAN"}, ', + b'{"mode": "NULLABLE", "name": "some_bit", "type": "INTEGER"}]', +] + + +class TestMsSqlToGoogleCloudStorageOperator: + @pytest.mark.parametrize( + "value, expected", + [ + ("string", "string"), + (32.9, 32.9), + (-2, -2), + (datetime.date(1970, 1, 2), "1970-01-02"), + (datetime.date(1000, 1, 2), "1000-01-02"), + (datetime.datetime(1970, 1, 1, 1, 0), "1970-01-01T01:00:00"), + (datetime.time(hour=0, minute=0, second=0), "00:00:00"), + (datetime.time(hour=23, minute=59, second=59), "23:59:59"), + ], + ) + def test_convert_type(self, value, expected): + op = MSSQLToGCSOperator( + task_id=TASK_ID, + mssql_conn_id=MSSQL_CONN_ID, + sql=SQL, + bucket=BUCKET, + filename=JSON_FILENAME, + ) + assert op.convert_type(value, None) == expected + + def test_init(self): + """Test MySqlToGoogleCloudStorageOperator instance is properly initialized.""" + op = MSSQLToGCSOperator(task_id=TASK_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME) + assert op.task_id == TASK_ID + assert op.sql == SQL + assert op.bucket == BUCKET + assert op.filename == JSON_FILENAME + + @mock.patch("airflow.providers.google.cloud.transfers.mssql_to_gcs.MsSqlHook") + @mock.patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + def test_exec_success_json(self, gcs_hook_mock_class, mssql_hook_mock_class): + """Test successful run of execute function for JSON""" + op = MSSQLToGCSOperator( + task_id=TASK_ID, mssql_conn_id=MSSQL_CONN_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME + ) + + mssql_hook_mock = mssql_hook_mock_class.return_value + mssql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) + mssql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION + + gcs_hook_mock = gcs_hook_mock_class.return_value + + def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False, metadata=None): + assert BUCKET == bucket + assert JSON_FILENAME.format(0) == obj + assert "application/json" == mime_type + assert GZIP == gzip + with open(tmp_filename, "rb") as file: + assert b"".join(NDJSON_LINES) == file.read() + + gcs_hook_mock.upload.side_effect = _assert_upload + + op.execute(None) + + mssql_hook_mock_class.assert_called_once_with(mssql_conn_id=MSSQL_CONN_ID) + mssql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL) + + @mock.patch("airflow.providers.google.cloud.transfers.mssql_to_gcs.MsSqlHook") + @mock.patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + def test_file_splitting(self, gcs_hook_mock_class, mssql_hook_mock_class): + """Test that ndjson is split by approx_max_file_size_bytes param.""" + mssql_hook_mock = mssql_hook_mock_class.return_value + mssql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) + mssql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION + + gcs_hook_mock = gcs_hook_mock_class.return_value + expected_upload = { + JSON_FILENAME.format(0): b"".join(NDJSON_LINES[:2]), + JSON_FILENAME.format(1): NDJSON_LINES[2], + } + + def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False, metadata=None): + assert BUCKET == bucket + assert "application/json" == mime_type + assert GZIP == gzip + with open(tmp_filename, "rb") as file: + assert expected_upload[obj] == file.read() + + gcs_hook_mock.upload.side_effect = _assert_upload + + op = MSSQLToGCSOperator( + task_id=TASK_ID, + sql=SQL, + bucket=BUCKET, + filename=JSON_FILENAME, + approx_max_file_size_bytes=len(expected_upload[JSON_FILENAME.format(0)]), + ) + op.execute(None) + + @mock.patch("airflow.providers.google.cloud.transfers.mssql_to_gcs.MsSqlHook") + @mock.patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + @pytest.mark.parametrize( + "bit_fields,schema_json", [(None, SCHEMA_JSON), (["bit_fields", SCHEMA_JSON_BIT_FIELDS])] + ) + def test_schema_file(self, gcs_hook_mock_class, mssql_hook_mock_class, bit_fields, schema_json): + """Test writing schema files.""" + mssql_hook_mock = mssql_hook_mock_class.return_value + mssql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) + mssql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION + + gcs_hook_mock = gcs_hook_mock_class.return_value + + def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip, metadata=None): + if obj == SCHEMA_FILENAME: + with open(tmp_filename, "rb") as file: + assert b"".join(SCHEMA_JSON) == file.read() + + gcs_hook_mock.upload.side_effect = _assert_upload + + op = MSSQLToGCSOperator( + task_id=TASK_ID, + sql=SQL, + bucket=BUCKET, + filename=JSON_FILENAME, + schema_filename=SCHEMA_FILENAME, + bit_fields=["some_bit"], + ) + op.execute(None) + + # once for the file and once for the schema + assert 2 == gcs_hook_mock.upload.call_count diff --git a/tests/providers/google/cloud/transfers/test_trino_to_gcs.py b/tests/providers/google/cloud/transfers/test_trino_to_gcs.py new file mode 100644 index 0000000000000..09da568f6a81b --- /dev/null +++ b/tests/providers/google/cloud/transfers/test_trino_to_gcs.py @@ -0,0 +1,327 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest.mock import patch + +from airflow.providers.google.cloud.transfers.trino_to_gcs import TrinoToGCSOperator + +TASK_ID = "test-trino-to-gcs" +TRINO_CONN_ID = "my-trino-conn" +GCP_CONN_ID = "my-gcp-conn" +IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"] +SQL = "SELECT * FROM memory.default.test_multiple_types" +BUCKET = "gs://test" +FILENAME = "test_{}.ndjson" + +NDJSON_LINES = [ + b'{"some_num": 42, "some_str": "mock_row_content_1"}\n', + b'{"some_num": 43, "some_str": "mock_row_content_2"}\n', + b'{"some_num": 44, "some_str": "mock_row_content_3"}\n', +] +CSV_LINES = [ + b"some_num,some_str\r\n", + b"42,mock_row_content_1\r\n", + b"43,mock_row_content_2\r\n", + b"44,mock_row_content_3\r\n", +] +SCHEMA_FILENAME = "schema_test.json" +SCHEMA_JSON = b'[{"name": "some_num", "type": "INT64"}, {"name": "some_str", "type": "STRING"}]' + + +class TestTrinoToGCSOperator: + def test_init(self): + """Test TrinoToGCSOperator instance is properly initialized.""" + op = TrinoToGCSOperator( + task_id=TASK_ID, + sql=SQL, + bucket=BUCKET, + filename=FILENAME, + impersonation_chain=IMPERSONATION_CHAIN, + ) + assert op.task_id == TASK_ID + assert op.sql == SQL + assert op.bucket == BUCKET + assert op.filename == FILENAME + assert op.impersonation_chain == IMPERSONATION_CHAIN + + @patch("airflow.providers.google.cloud.transfers.trino_to_gcs.TrinoHook") + @patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + def test_save_as_json(self, mock_gcs_hook, mock_trino_hook): + def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip, metadata=None): + assert BUCKET == bucket + assert FILENAME.format(0) == obj + assert "application/json" == mime_type + assert not gzip + with open(tmp_filename, "rb") as file: + assert b"".join(NDJSON_LINES) == file.read() + + mock_gcs_hook.return_value.upload.side_effect = _assert_upload + + mock_cursor = mock_trino_hook.return_value.get_conn.return_value.cursor + + mock_cursor.return_value.description = [ + ("some_num", "INTEGER", None, None, None, None, None), + ("some_str", "VARCHAR", None, None, None, None, None), + ] + + mock_cursor.return_value.fetchone.side_effect = [ + [42, "mock_row_content_1"], + [43, "mock_row_content_2"], + [44, "mock_row_content_3"], + None, + ] + + op = TrinoToGCSOperator( + task_id=TASK_ID, + sql=SQL, + bucket=BUCKET, + filename=FILENAME, + trino_conn_id=TRINO_CONN_ID, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + ) + + op.execute(None) + + mock_trino_hook.assert_called_once_with(trino_conn_id=TRINO_CONN_ID) + mock_gcs_hook.assert_called_once_with( + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + ) + + mock_gcs_hook.return_value.upload.assert_called() + + @patch("airflow.providers.google.cloud.transfers.trino_to_gcs.TrinoHook") + @patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + def test_save_as_json_with_file_splitting(self, mock_gcs_hook, mock_trino_hook): + """Test that ndjson is split by approx_max_file_size_bytes param.""" + + expected_upload = { + FILENAME.format(0): b"".join(NDJSON_LINES[:2]), + FILENAME.format(1): NDJSON_LINES[2], + } + + def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip, metadata=None): + assert BUCKET == bucket + assert "application/json" == mime_type + assert not gzip + with open(tmp_filename, "rb") as file: + assert expected_upload[obj] == file.read() + + mock_gcs_hook.return_value.upload.side_effect = _assert_upload + + mock_cursor = mock_trino_hook.return_value.get_conn.return_value.cursor + + mock_cursor.return_value.description = [ + ("some_num", "INTEGER", None, None, None, None, None), + ("some_str", "VARCHAR(20)", None, None, None, None, None), + ] + + mock_cursor.return_value.fetchone.side_effect = [ + [42, "mock_row_content_1"], + [43, "mock_row_content_2"], + [44, "mock_row_content_3"], + None, + ] + + op = TrinoToGCSOperator( + task_id=TASK_ID, + sql=SQL, + bucket=BUCKET, + filename=FILENAME, + approx_max_file_size_bytes=len(expected_upload[FILENAME.format(0)]), + ) + + op.execute(None) + + mock_gcs_hook.return_value.upload.assert_called() + + @patch("airflow.providers.google.cloud.transfers.trino_to_gcs.TrinoHook") + @patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + def test_save_as_json_with_schema_file(self, mock_gcs_hook, mock_trino_hook): + """Test writing schema files.""" + + def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip, metadata=None): + if obj == SCHEMA_FILENAME: + with open(tmp_filename, "rb") as file: + assert SCHEMA_JSON == file.read() + + mock_gcs_hook.return_value.upload.side_effect = _assert_upload + + mock_cursor = mock_trino_hook.return_value.get_conn.return_value.cursor + + mock_cursor.return_value.description = [ + ("some_num", "INTEGER", None, None, None, None, None), + ("some_str", "VARCHAR", None, None, None, None, None), + ] + + mock_cursor.return_value.fetchone.side_effect = [ + [42, "mock_row_content_1"], + [43, "mock_row_content_2"], + [44, "mock_row_content_3"], + None, + ] + + op = TrinoToGCSOperator( + task_id=TASK_ID, + sql=SQL, + bucket=BUCKET, + filename=FILENAME, + schema_filename=SCHEMA_FILENAME, + export_format="csv", + trino_conn_id=TRINO_CONN_ID, + gcp_conn_id=GCP_CONN_ID, + ) + op.execute(None) + + # once for the file and once for the schema + assert 2 == mock_gcs_hook.return_value.upload.call_count + + @patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + @patch("airflow.providers.google.cloud.transfers.trino_to_gcs.TrinoHook") + def test_save_as_csv(self, mock_trino_hook, mock_gcs_hook): + def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip, metadata=None): + assert BUCKET == bucket + assert FILENAME.format(0) == obj + assert "text/csv" == mime_type + assert not gzip + with open(tmp_filename, "rb") as file: + assert b"".join(CSV_LINES) == file.read() + + mock_gcs_hook.return_value.upload.side_effect = _assert_upload + + mock_cursor = mock_trino_hook.return_value.get_conn.return_value.cursor + + mock_cursor.return_value.description = [ + ("some_num", "INTEGER", None, None, None, None, None), + ("some_str", "VARCHAR", None, None, None, None, None), + ] + + mock_cursor.return_value.fetchone.side_effect = [ + [42, "mock_row_content_1"], + [43, "mock_row_content_2"], + [44, "mock_row_content_3"], + None, + ] + + op = TrinoToGCSOperator( + task_id=TASK_ID, + sql=SQL, + bucket=BUCKET, + filename=FILENAME, + export_format="csv", + trino_conn_id=TRINO_CONN_ID, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + ) + + op.execute(None) + + mock_gcs_hook.return_value.upload.assert_called() + + mock_trino_hook.assert_called_once_with(trino_conn_id=TRINO_CONN_ID) + mock_gcs_hook.assert_called_once_with( + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + ) + + @patch("airflow.providers.google.cloud.transfers.trino_to_gcs.TrinoHook") + @patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + def test_save_as_csv_with_file_splitting(self, mock_gcs_hook, mock_trino_hook): + """Test that csv is split by approx_max_file_size_bytes param.""" + + expected_upload = { + FILENAME.format(0): b"".join(CSV_LINES[:3]), + FILENAME.format(1): b"".join([CSV_LINES[0], CSV_LINES[3]]), + } + + def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip, metadata=None): + assert BUCKET == bucket + assert "text/csv" == mime_type + assert not gzip + with open(tmp_filename, "rb") as file: + assert expected_upload[obj] == file.read() + + mock_gcs_hook.return_value.upload.side_effect = _assert_upload + + mock_cursor = mock_trino_hook.return_value.get_conn.return_value.cursor + + mock_cursor.return_value.description = [ + ("some_num", "INTEGER", None, None, None, None, None), + ("some_str", "VARCHAR(20)", None, None, None, None, None), + ] + + mock_cursor.return_value.fetchone.side_effect = [ + [42, "mock_row_content_1"], + [43, "mock_row_content_2"], + [44, "mock_row_content_3"], + None, + ] + + op = TrinoToGCSOperator( + task_id=TASK_ID, + sql=SQL, + bucket=BUCKET, + filename=FILENAME, + approx_max_file_size_bytes=len(expected_upload[FILENAME.format(0)]), + export_format="csv", + ) + + op.execute(None) + + mock_gcs_hook.return_value.upload.assert_called() + + @patch("airflow.providers.google.cloud.transfers.trino_to_gcs.TrinoHook") + @patch("airflow.providers.google.cloud.transfers.sql_to_gcs.GCSHook") + def test_save_as_csv_with_schema_file(self, mock_gcs_hook, mock_trino_hook): + """Test writing schema files.""" + + def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip, metadata=None): + if obj == SCHEMA_FILENAME: + with open(tmp_filename, "rb") as file: + assert SCHEMA_JSON == file.read() + + mock_gcs_hook.return_value.upload.side_effect = _assert_upload + + mock_cursor = mock_trino_hook.return_value.get_conn.return_value.cursor + + mock_cursor.return_value.description = [ + ("some_num", "INTEGER", None, None, None, None, None), + ("some_str", "VARCHAR", None, None, None, None, None), + ] + + mock_cursor.return_value.fetchone.side_effect = [ + [42, "mock_row_content_1"], + [43, "mock_row_content_2"], + [44, "mock_row_content_3"], + None, + ] + + op = TrinoToGCSOperator( + task_id=TASK_ID, + sql=SQL, + bucket=BUCKET, + filename=FILENAME, + schema_filename=SCHEMA_FILENAME, + export_format="csv", + ) + op.execute(None) + + # once for the file and once for the schema + assert 2 == mock_gcs_hook.return_value.upload.call_count