feat: implement testing framework for postgres

.pre-commit-config.yaml

@@ -32,7 +32,7 @@ repos:
 
   # Ruff replaces black, flake8, autoflake and isort
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.7.1" # make sure this is always consistent with hatch configs
+    rev: "v0.9.3" # make sure this is always consistent with hatch configs
     hooks:
       - id: ruff
         args: [--config, ./pyproject.toml]

.vscode/settings.json

@@ -150,8 +150,7 @@
     "source.convertImportFormat"
   ],
   "sqltools.disableReleaseNotifications": true,
-  "sqltools.disableNodeDetectNotifications": true,
-  "cloudcode.duetAI.enable": true,
+"sqltools.disableNodeDetectNotifications": true,
   "cloudcode.compute.sshInternalIp": true,
   "python.testing.pytestArgs": [
       "tests"
@@ -187,5 +186,6 @@
 "python.analysis.experimentalserver": true,
 "python.analysis.diagnosticSeverityOverrides": {
     "reportUnknownMemberType": "none"
-}
+},
+"geminicodeassist.enable": true
 }

Makefile

@@ -31,23 +31,20 @@ install-pipx: 										## Install pipx
 	@python3 -m pip install --upgrade --user pipx
 
 install-hatch: 										## Install Hatch, UV, and Ruff
-	@pipx install hatch --force
-	@pipx inject hatch ruff uv hatch-pip-compile hatch-vcs --include-deps --include-apps --force
+	@sh ./tools/install-hatch.sh
 
 configure-hatch: 										## Configure Hatch defaults
 	@hatch config set dirs.env.virtual .direnv
 	@hatch config set dirs.env.pip-compile .direnv
 
 upgrade-hatch: 										## Update Hatch, UV, and Ruff
-	@pipx upgrade hatch --include-injected
+	@hatch self update
 
 install: 										## Install the project and all dependencies
 	@if [ "$(VENV_EXISTS)" ]; then echo "=> Removing existing virtual environment"; $(MAKE) destroy-venv; fi
 	@$(MAKE) clean
 	@if [ "$(NODE_MODULES_EXISTS)" ]; then echo "=> Removing existing node modules"; $(MAKE) destroy-node_modules; fi
-	@if ! pipx --version > /dev/null; then echo '=> Installing `pipx`'; $(MAKE) install-pipx ; fi
-	@if ! hatch --version > /dev/null; then echo '=> Installing `hatch` with `pipx`'; $(MAKE) install-hatch ; fi
-	@if ! hatch-pip-compile --version > /dev/null; then echo '=> Updating `hatch` and installing plugins'; $(MAKE) upgrade-hatch ; fi
+	@if ! hatch --version > /dev/null; then echo '=> Installing `hatch`'; $(MAKE) install-hatch ; fi
 	@echo "=> Creating Python environments..."
 	@$(MAKE) configure-hatch
 	@hatch env create local
@@ -58,12 +55,12 @@ install: 										## Install the project and all dependencies
 .PHONY: upgrade
 upgrade:       										## Upgrade all dependencies to the latest stable versions
 	@echo "=> Updating all dependencies"
-	@hatch-pip-compile --upgrade --all
 	@echo "=> Python Dependencies Updated"
 	@if [ "$(USING_NPM)" ]; then hatch run local:npm upgrade --latest; fi
 	@echo "=> Node Dependencies Updated"
 	@hatch run lint:pre-commit autoupdate
 	@echo "=> Updated Pre-commit"
+	@$(MAKE) install
 
 
 .PHONY: clean
@@ -245,8 +242,8 @@ test:  												## Run the tests
 	@hatch run +py="3.12" test:cov
 	@echo "=> Tests complete"
 
-.PHONY: test-all
-test-all:  												## Run the tests against all python versions
+.PHONY: test-all-pythons
+test-all-pythons:  												## Run the tests against all python versions
 	@echo "=> Running test cases"
 	@hatch run test:cov
 	@echo "=> Tests complete"

docs/user_guide/sqlserver/collection_scripts.md → docs/user_guide/shell_scripts/sqlserver/collection_scripts.md

@@ -311,12 +311,12 @@ Occasionally, organizational security policies require that Powershell scripts b
             DnsName = '[email protected]'
             FriendlyName = 'Google DMA Self Signed PS Code Signing'
             NotAfter = (Get-Date).AddYears(5)
-            Type = 'CodeSigning'  
-            CertStoreLocation = 'cert:\CurrentUser\My' 
+            Type = 'CodeSigning'
+            CertStoreLocation = 'cert:\CurrentUser\My'
             KeyUsage = 'DigitalSignature'
             KeyAlgorithm = 'RSA'
-            KeyLength = 2048 
-            HashAlgorithm = 'sha256' 
+            KeyLength = 2048
+            HashAlgorithm = 'sha256'
             }
 
     - Create a new self-signed certificate based on the above parameters and send the details to 'newCodeSigningCert' variable for reference later.

mkdocs.yml

@@ -16,22 +16,24 @@ plugins:
 
 nav:
   - Home: index.md
-  - Oracle Data Collection:
-      - Data Collection: user_guide/oracle/collection_scripts.md
-      - Permissions Required: user_guide/oracle/permissions.md
-      - Database User Scripts: user_guide/oracle/db_user_create.md
-      - Collection Masking: user_guide/oracle/collection_masker.md
-  - Microsoft SQL Server Data Collection:
-      - Data Collection: user_guide/sqlserver/collection_scripts.md
-      - Database User Scripts: user_guide/sqlserver/db_user_create.md
-  - Postgres Data Collection:
-      - Data Collection: user_guide/postgres/collection_scripts.md
-  - MySQL Data Collection:
-      - Data Collection: user_guide/mysql/collection_scripts.md
-  - Developers:
-      - Developer Setup: developer_guide/developer_setup.md
-      - Commands: developer_guide/commands.md
-      - Workflows: developer_guide/workflows.md
-      - Releases: developer_guide/releases.md
+  - Readiness Check: user_guide/readiness_check/index.md
+  - Shell Script Collectors:
+    - Oracle Data Collection:
+        - Data Collection: user_guide/shell_scripts/oracle/collection_scripts.md
+        - Permissions Required: user_guide/shell_scripts/oracle/permissions.md
+        - Database User Scripts: user_guide/shell_scripts/oracle/db_user_create.md
+        - Collection Masking: user_guide/shell_scripts/oracle/collection_masker.md
+    - Microsoft SQL Server Data Collection:
+        - Data Collection: user_guide/shell_scripts/sqlserver/collection_scripts.md
+        - Database User Scripts: user_guide/shell_scripts/sqlserver/db_user_create.md
+    - Postgres Data Collection:
+        - Data Collection: user_guide/shell_scripts/postgres/collection_scripts.md
+    - MySQL Data Collection:
+        - Data Collection: user_guide/shell_scripts/mysql/collection_scripts.md
+    - Developers:
+        - Developer Setup: developer_guide/developer_setup.md
+        - Commands: developer_guide/commands.md
+        - Workflows: developer_guide/workflows.md
+        - Releases: developer_guide/releases.md
   - Code Of Conduct: code_of_conduct.md
   - Contributing: contributing.md

pyproject.toml

@@ -62,7 +62,7 @@ dependencies = [
     "sqlalchemy>=2.0.25",
     "typing-extensions>=4.0.0",
     "msgspec",
-    "greenlet; sys_platform == \"darwin\"",
+    "greenlet",
 ]
 
 [project.urls]
@@ -75,7 +75,7 @@ Source = "https://github.com/GoogleCloudPlatform/database-assessment"
 mssql = ["aioodbc"]
 mysql = ["asyncmy>=0.2.9"]
 oracle = ["oracledb"]
-postgres = ["asyncpg>=0.29.0"]
+postgres = ["psycopg[pool,binary]"]
 server = ["litestar[structlog,jinja]>=2.7.0", "litestar-granian>=0.2.3"]
 
 
@@ -128,7 +128,7 @@ extra-dependencies = [
     "pytest-sugar",
     "pytest-click",
     "pytest-xdist",
-    "pytest-databases",
+    "pytest-databases<=0.10.0",
     # lint
     "mypy",
     "ruff",

scripts/collector/postgres/collect-data.sh

@@ -469,21 +469,6 @@ DIAGPACKACCESS="postgres"
 
  fi
 
-#############################################################################
-
-# if [[  $# -lt 3  || $# -gt 4 || (  "$2" != "UseDiagnostics" && "$2" != "NoDiagnostics" ) ]]
-#  then
-#   echo
-#   echo "You must indicate whether or not to use the Diagnostics Pack views."
-#   echo "If this database is licensed to use the Diagnostics pack:"
-#   echo "  $0 $1 UseDiagnostics"
-#   echo " "
-#   echo "If this database is NOT licensed to use the Diagnostics pack:"
-#   echo "  $0 $1 NoDiagnostics"
-#   echo " "
-#   exit 1
-# fi
-
 # MAIN
 #############################################################################
 

scripts/collector/sqlserver/README.txt

@@ -286,12 +286,12 @@ Operating System Versions:
                 DnsName = '[email protected]'
                 FriendlyName = 'Google DMA Self Signed PS Code Signing'
                 NotAfter = (Get-Date).AddYears(5)
-                Type = 'CodeSigning'  
-                CertStoreLocation = 'cert:\CurrentUser\My' 
+                Type = 'CodeSigning'
+                CertStoreLocation = 'cert:\CurrentUser\My'
                 KeyUsage = 'DigitalSignature'
                 KeyAlgorithm = 'RSA'
-                KeyLength = 2048 
-                HashAlgorithm = 'sha256' 
+                KeyLength = 2048
+                HashAlgorithm = 'sha256'
                 }
 
         - Create a new self-signed certificate based on the above parameters and send the details to 'newCodeSigningCert' variable for reference later.

scripts/collector/sqlserver/instanceReview.ps1

@@ -726,4 +726,4 @@ else {
     WriteLog -logLocation $foldername\$logFile -logMessage "Collection Complete..." -logOperation "MESSAGE"
 }
 
-Exit 0
+Exit 0

scripts/collector/sqlserver/sql/createCollectionUser.sql

@@ -183,4 +183,4 @@ BEGIN
 	FETCH NEXT FROM db_cursor INTO @dbname;
 	CLOSE db_cursor
 	DEALLOCATE db_cursor
-END;
+END;

scripts/collector/sqlserver/sql/dbServerFeatures.sql

@@ -170,7 +170,7 @@ BEGIN
                         ELSE 0
                     END AS Is_EnabledOrUsed,
                     dqs_count as Count
-                from dqs_service');  
+                from dqs_service');
         END CATCH
 END;
 

src/dma/__init__.py

@@ -30,7 +30,7 @@
 click.rich_click.ERRORS_EPILOGUE = """
 For additional support, refer to the documentation at https://googlecloudplatform.github.io/database-assessment/
 """
-click.rich_click.MAX_WIDTH = 80
+click.rich_click.WIDTH = 80
 click.rich_click.SHOW_METAVARS_COLUMN = True
 click.rich_click.APPEND_METAVARS_HELP = True
 click.rich_click.STYLE_OPTION = "bold cyan"

src/dma/cli/main.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 from __future__ import annotations
 
-import asyncio
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal
@@ -147,25 +146,23 @@ def collect_data(
         password = prompt.Prompt.ask("Please enter a password", password=True)
     input_confirmed = True if no_prompt else prompt.Confirm.ask("Are you ready to start the assessment?")
     if input_confirmed:
-        asyncio.run(
-            _collect_data(
-                console=console,
-                src_info=SourceInfo(
-                    db_type=db_type.upper(),  # type: ignore[arg-type]
-                    username=username,
-                    password=password,
-                    hostname=hostname,
-                    port=port,
-                ),
-                database=database,
-                collection_identifier=collection_identifier,
-            )
+        _collect_data(
+            console=console,
+            src_info=SourceInfo(
+                db_type=db_type.upper(),  # type: ignore[arg-type]
+                username=username,
+                password=password,
+                hostname=hostname,
+                port=port,
+            ),
+            database=database,
+            collection_identifier=collection_identifier,
         )
     else:
         console.rule("Skipping execution until input is confirmed", align="left")
 
 
-async def _collect_data(
+def _collect_data(
     console: Console,
     src_info: SourceInfo,
     database: str,
@@ -184,7 +181,7 @@ async def _collect_data(
             console=console,
             collection_identifier=collection_identifier,
         )
-        await collection_extractor.execute()
+        collection_extractor.execute()
         collection_extractor.dump_database(working_path)
 
 
@@ -290,25 +287,23 @@ def readiness_assessment(
         password = prompt.Prompt.ask("Please enter a password", password=True)
     input_confirmed = True if no_prompt else prompt.Confirm.ask("Are you ready to start the assessment?")
     if input_confirmed:
-        asyncio.run(
-            _readiness_check(
-                console=console,
-                src_info=SourceInfo(
-                    db_type=db_type.upper(),  # type: ignore[arg-type]
-                    username=username,
-                    password=password,
-                    hostname=hostname,
-                    port=port,
-                ),
-                database=database,
-                collection_identifier=collection_identifier,
-            )
+        _readiness_check(
+            console=console,
+            src_info=SourceInfo(
+                db_type=db_type.upper(),  # type: ignore[arg-type]
+                username=username,
+                password=password,
+                hostname=hostname,
+                port=port,
+            ),
+            database=database,
+            collection_identifier=collection_identifier,
         )
     else:
         console.rule("Skipping execution until input is confirmed", align="left")
 
 
-async def _readiness_check(
+def _readiness_check(
     console: Console,
     src_info: SourceInfo,
     database: str,
@@ -326,7 +321,7 @@ async def _readiness_check(
             collection_identifier=collection_identifier,
             working_path=working_path,
         )
-        await workflow.execute()
+        workflow.execute()
         console.print(Padding("", 1, expand=True))
         console.rule("Processing collected data.", align="left")
         workflow.print_summary()

src/dma/collector/dependencies.py

@@ -27,34 +27,37 @@
 from dma.lib.exceptions import ApplicationError
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator, Generator
+    from collections.abc import Generator, Iterator
     from pathlib import Path
 
     import duckdb
-    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
 
 
-async def provide_collection_query_manager(
-    db_session: AsyncSession,
+def provide_collection_query_manager(
+    db_session: Session,
     execution_id: str | None = None,
     source_id: str | None = None,
     manual_id: str | None = None,
-) -> AsyncIterator[CollectionQueryManager]:
+) -> Iterator[CollectionQueryManager]:
     """Provide collection query manager.
 
     Uses SQLAlchemy Connection management to establish and retrieve a valid database session.
 
     The driver dialect is detected from the session and the underlying raw DBAPI connection is fetched and passed to the Query Manager.
     """
     dialect = db_session.bind.dialect if db_session.bind is not None else db_session.get_bind().dialect
-    db_connection = await db_session.connection()
+    db_connection = db_session.connection()
 
-    raw_connection = await db_connection.get_raw_connection()
+    raw_connection = db_connection.engine.raw_connection()
     if not raw_connection.driver_connection:
         msg = "Unable to fetch raw connection from session."
         raise ApplicationError(msg)
     rdbms_type = dialect.name
     if rdbms_type == "postgresql":
+        from psycopg.rows import dict_row  # noqa: PLC0415
+
+        raw_connection.driver_connection.row_factory = dict_row
         query_manager: CollectionQueryManager = PostgresCollectionQueryManager(
             connection=raw_connection.driver_connection,
             manual_id=manual_id,