├── .envrc.example ├── .flake8 ├── .github ├── CODEOWNERS ├── ISSUE_TEMPLATE │ ├── bug_report.yaml │ ├── documentation_issue.yaml │ ├── feature_request.yaml │ └── question.yaml ├── auto-label.yml ├── blunderbuss.yml ├── flakybot.yaml ├── header-checker-lint.yml ├── labels.yml ├── release-please.yml ├── release-trigger.yml ├── renovate.json5 └── workflows │ ├── codeql.yaml │ ├── coverage.yaml │ ├── labels.yaml │ ├── lint.yaml │ ├── scorecard.yaml │ └── tests.yaml ├── .gitignore ├── CHANGELOG.md ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── SECURITY.md ├── docs └── images │ └── alloydb-python-connector.png ├── google ├── api │ └── field_behavior_pb2.py └── cloud │ ├── alloydb │ └── connector │ │ └── __init__.py │ ├── alloydb_connectors_v1 │ └── proto │ │ ├── __init__.py │ │ ├── resources_pb2.py │ │ └── resources_pb2.pyi │ └── alloydbconnector │ ├── __init__.py │ ├── async_connector.py │ ├── asyncpg.py │ ├── client.py │ ├── connection_info.py │ ├── connector.py │ ├── enums.py │ ├── exceptions.py │ ├── instance.py │ ├── lazy.py │ ├── pg8000.py │ ├── py.typed │ ├── rate_limiter.py │ ├── refresh_utils.py │ ├── static.py │ ├── types.py │ ├── utils.py │ └── version.py ├── noxfile.py ├── pyproject.toml ├── requirements-test.txt ├── requirements.txt └── tests ├── system ├── test_alloydb_connector_package.py ├── test_asyncpg_connection.py ├── test_asyncpg_iam_authn.py ├── test_asyncpg_psc.py ├── test_asyncpg_public_ip.py ├── test_native_asyncpg_direct_connection.py ├── test_pg8000_connection.py ├── test_pg8000_iam_authn.py ├── test_pg8000_psc.py ├── test_pg8000_public_ip.py ├── test_psycopg2_direct_connection.py └── test_sqlalchemy_asyncpg_direct_connection.py └── unit ├── conftest.py ├── mocks.py ├── test_async_connector.py ├── test_client.py ├── test_connection_info.py ├── test_connector.py ├── test_instance.py ├── test_lazy.py ├── test_packaging.py ├── test_rate_limiter.py ├── test_refresh_utils.py ├── test_static.py └── test_utils.py /.envrc.example: -------------------------------------------------------------------------------- 1 | export ALLOYDB_DB="some-db" 2 | export ALLOYDB_USER="some-user" 3 | export ALLOYDB_PASS="some-password" 4 | export ALLOYDB_INSTANCE_NAME="projects//locations//clusters//instances/" 5 | export ALLOYDB_INSTANCE_IP="some-IP-address" 6 | export ALLOYDB_IAM_USER="some-user@my-project.iam" 7 | export ALLOYDB_IMPERSONATED_USER="some-impersonated-IAM-user" 8 | export ALLOYDB_PSC_INSTANCE_URI="projects//locations//clusters//instances/" 9 | -------------------------------------------------------------------------------- /.flake8: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # 3 | # Copyright 2023 Google LLC 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # https://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | 17 | [flake8] 18 | ignore = E203, E231, E266, E501, W503, ANN101, ANN102, ANN401 19 | exclude = 20 | # Exclude generated code. 21 | **/proto/** 22 | **/gapic/** 23 | **/services/** 24 | **/types/** 25 | *_pb2.py 26 | 27 | # Standard linting exemptions. 28 | **/.nox/** 29 | __pycache__, 30 | .git, 31 | *.pyc, 32 | conf.py 33 | -------------------------------------------------------------------------------- /.github/CODEOWNERS: -------------------------------------------------------------------------------- 1 | * @GoogleCloudPlatform/alloydb-connectors-code-owners 2 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug_report.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: 🐞 Bug Report 16 | description: File a bug report 17 | title: "Brief summary of what bug or error was observed" 18 | labels: ["type: bug"] 19 | body: 20 | - type: markdown 21 | attributes: 22 | value: | 23 | Thanks for stopping by to let us know something could be better! 24 | 25 | Please run down the following list and make sure you've tried the usual "quick fixes": 26 | - Search the [current open issues](https://github.com/GoogleCloudPlatform/alloydb-python-connector/issues) 27 | - Check for answers on [StackOverflow](https://stackoverflow.com/questions/tagged/google-alloydb) (under the 'google-alloydb' tag) 28 | 29 | If you are still having issues, please include as much information as possible below! :smile: 30 | - type: textarea 31 | id: bug-description 32 | attributes: 33 | label: Bug Description 34 | description: "Please enter a detailed description of the bug, and any information about what behavior you noticed and why it is defective or unintentional." 35 | validations: 36 | required: true 37 | - type: textarea 38 | id: example-code 39 | attributes: 40 | label: Example code (or command) 41 | description: "Please paste any useful application code related to the bug below. (if your code is in a public repo, feel free to paste a link!)" 42 | render: Python 43 | - type: textarea 44 | id: stacktrace 45 | attributes: 46 | label: Stacktrace 47 | description: "Paste any relevant stacktrace or error you are running into here. Be sure to filter sensitive information!" 48 | render: bash 49 | - type: textarea 50 | id: repro 51 | attributes: 52 | label: Steps to reproduce? 53 | description: "How do you trigger this bug? Please walk us through it step by step." 54 | value: | 55 | 1. ? 56 | 2. ? 57 | 3. ? 58 | ... 59 | validations: 60 | required: true 61 | - type: textarea 62 | id: environment 63 | attributes: 64 | label: Environment 65 | description: "Let us know some details about the environment in which you are seeing the bug!" 66 | value: | 67 | 1. OS type and version: 68 | 2. Python version: 69 | 3. AlloyDB Python Connector version: 70 | validations: 71 | required: true 72 | - type: textarea 73 | id: additional-details 74 | attributes: 75 | label: Additional Details 76 | description: "Any other information you want us to know?" 77 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/documentation_issue.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: 📝 Documentation Issue 16 | description: Report wrong or missing information with the documentation in this repo. 17 | title: "Brief summary of what is missing or incorrect" 18 | labels: ["type: docs"] 19 | body: 20 | - type: markdown 21 | attributes: 22 | value: | 23 | Thanks for stopping by to let us know something could be better! :smile: 24 | 25 | Please explain below how we can improve our documentation. 26 | - type: textarea 27 | id: description 28 | attributes: 29 | label: Description 30 | description: "Provide a short description of what is missing or incorrect, as well as a link to the specific location of the issue." 31 | validations: 32 | required: true 33 | - type: textarea 34 | id: potential-solution 35 | attributes: 36 | label: Potential Solution 37 | description: "What would you prefer the documentation say? Why would this information be more accurate or helpful?" 38 | - type: textarea 39 | id: additional-details 40 | attributes: 41 | label: Additional Details 42 | description: "Please reference any other relevant issues, PRs, descriptions, or screenshots here." 43 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: ✨ Feature Request 16 | description: Suggest an idea for new or improved behavior. 17 | title: "Brief summary of the proposed feature" 18 | labels: ["type: feature request"] 19 | body: 20 | - type: markdown 21 | attributes: 22 | value: | 23 | Thanks for stopping by to let us know something could be better! 24 | 25 | Please run down the following list before proceeding with your feature request: 26 | - Search the [current open issues](https://github.com/GoogleCloudPlatform/alloydb-python-connector/issues) to prevent creating a duplicate. 27 | 28 | Please include as much information as possible below! :smile: 29 | - type: textarea 30 | id: feature-description 31 | attributes: 32 | label: Feature Description 33 | description: "A clear and concise description of what feature you would like to see, and why it would be useful to have added." 34 | validations: 35 | required: true 36 | - type: textarea 37 | id: sample-code 38 | attributes: 39 | label: Sample code 40 | description: "If you already have an idea of what the implementation of this feature would like in code please provide it. (pseudo code is okay!)" 41 | render: Python 42 | - type: textarea 43 | id: alternatives-considered 44 | attributes: 45 | label: Alternatives Considered 46 | description: "Are there any workaround or third party tools to replicate this behavior? Why would adding this feature be preferred over them?" 47 | - type: textarea 48 | id: additional-details 49 | attributes: 50 | label: Additional Details 51 | description: "Any additional information we should know? Please reference it here (issues, PRs, descriptions, or screenshots)" 52 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/question.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: 💬 Question 16 | description: Questions on how something works or the best way to do something? 17 | title: "Brief summary of your question" 18 | labels: ["type: question"] 19 | body: 20 | - type: markdown 21 | attributes: 22 | value: | 23 | Thanks for stopping by to let us know something could be better! 24 | 25 | Please run down the following list and make sure you've tried the usual "quick fixes": 26 | - Search the [current open issues](https://github.com/GoogleCloudPlatform/alloydb-python-connector/issues) for a similar question 27 | - Check for answers on [StackOverflow](https://stackoverflow.com/questions/tagged/google-alloydb) (under the 'google-alloydb' tag) 28 | 29 | If you still have a question, please include as much information as possible below! :smile: 30 | - type: textarea 31 | id: question 32 | attributes: 33 | label: Question 34 | description: "What's your question? Please provide as much relevant information as possible to reduce turnaround time." 35 | placeholder: "Example: How do I connect using this connector with Private IP from Cloud Run?" 36 | validations: 37 | required: true 38 | - type: textarea 39 | id: code 40 | attributes: 41 | label: Code 42 | description: "Please paste any useful application code that might be relevant to your question. (if your code is in a public repo, feel free to paste a link!)" 43 | render: Python 44 | - type: textarea 45 | id: additional-details 46 | attributes: 47 | label: Additional Details 48 | description: "Any other information you want us to know that might be helpful in answering your question? (link issues, PRs, descriptions, or screenshots)" 49 | -------------------------------------------------------------------------------- /.github/auto-label.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | enabled: false 16 | -------------------------------------------------------------------------------- /.github/blunderbuss.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | assign_issues: 16 | - enocom 17 | - nancynh 18 | - rhatgadkar-goog 19 | 20 | assign_prs: 21 | - enocom 22 | - nancynh 23 | - rhatgadkar-goog 24 | -------------------------------------------------------------------------------- /.github/flakybot.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | issuePriority: p2 16 | -------------------------------------------------------------------------------- /.github/header-checker-lint.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | allowedCopyrightHolders: 16 | - 'Google LLC' 17 | allowedLicenses: 18 | - 'Apache-2.0' 19 | sourceFileExtensions: 20 | - 'py' 21 | - 'yaml' 22 | - 'yml' 23 | - 'sh' 24 | -------------------------------------------------------------------------------- /.github/labels.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | - name: duplicate 16 | color: ededed 17 | description: "" 18 | 19 | - name: 'type: bug' 20 | color: db4437 21 | description: Error or flaw in code with unintended results or allowing sub-optimal 22 | usage patterns. 23 | - name: 'type: cleanup' 24 | color: c5def5 25 | description: An internal cleanup or hygiene concern. 26 | - name: 'type: docs' 27 | color: 0000A0 28 | description: Improvement to the documentation for an API. 29 | - name: 'type: feature request' 30 | color: c5def5 31 | description: ‘Nice-to-have’ improvement, new feature or different behavior or design. 32 | - name: 'type: process' 33 | color: c5def5 34 | description: A process-related concern. May include testing, release, or the like. 35 | - name: 'type: question' 36 | color: c5def5 37 | description: Request for information or clarification. 38 | 39 | - name: 'priority: p0' 40 | color: b60205 41 | description: Highest priority. Critical issue. P0 implies highest priority. 42 | - name: 'priority: p1' 43 | color: ffa03e 44 | description: Important issue which blocks shipping the next release. Will be fixed 45 | prior to next release. 46 | - name: 'priority: p2' 47 | color: fef2c0 48 | description: Moderately-important priority. Fix may not be included in next release. 49 | - name: 'priority: p3' 50 | color: ffffc7 51 | description: Desirable enhancement or fix. May not be included in next release. 52 | 53 | - name: automerge 54 | color: 00ff00 55 | description: Merge the pull request once unit tests and other checks pass. 56 | - name: 'automerge: exact' 57 | color: 8dd517 58 | description: Summon MOG for automerging, but approvals need to be against the latest 59 | commit 60 | - name: do not merge 61 | color: d93f0b 62 | description: Indicates a pull request not ready for merge, due to either quality 63 | or timing. 64 | 65 | - name: 'autorelease: pending' 66 | color: ededed 67 | description: Release please needs to do its work on this. 68 | - name: 'autorelease: triggered' 69 | color: ededed 70 | description: Release please has triggered a release for this. 71 | - name: 'autorelease: tagged' 72 | color: ededed 73 | description: Release please has completed a release for this. 74 | 75 | - name: 'flakybot: flaky' 76 | color: 86d9d7 77 | description: Tells the Flaky Bot not to close or comment on this issue. 78 | - name: 'flakybot: quiet' 79 | color: 86d9d7 80 | description: Tells the Flaky Bot to comment less. 81 | - name: 'flakybot: issue' 82 | color: a9f9f7 83 | description: An issue filed by the Flaky Bot. Should not be added manually. 84 | -------------------------------------------------------------------------------- /.github/release-please.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | handleGHRelease: true 16 | releaseType: python 17 | -------------------------------------------------------------------------------- /.github/release-trigger.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | enabled: true 16 | multiScmName: alloydb-python-connector 17 | -------------------------------------------------------------------------------- /.github/renovate.json5: -------------------------------------------------------------------------------- 1 | { 2 | "extends": [ 3 | "config:base", // https://docs.renovatebot.com/presets-config/#configbase 4 | ":semanticCommits", // https://docs.renovatebot.com/presets-default/#semanticcommits 5 | ":ignoreUnstable", // https://docs.renovatebot.com/presets-default/#ignoreunstable 6 | "group:allNonMajor", // https://docs.renovatebot.com/presets-group/#groupallnonmajor 7 | ":separateMajorReleases", // https://docs.renovatebot.com/presets-default/#separatemajorreleases 8 | ":prConcurrentLimitNone", // View complete backlog as PRs. https://docs.renovatebot.com/presets-default/#prconcurrentlimitnone 9 | ], 10 | "rebaseWhen": "behind-base-branch", 11 | "dependencyDashboard": true, 12 | "dependencyDashboardLabels": ["type: process"], 13 | "pip_requirements": { 14 | "fileMatch": ["requirements-test.txt"] 15 | }, 16 | "pip_setup": { 17 | "fileMatch": ["(^|/)setup\\.py$"] 18 | }, 19 | "packageRules": [ 20 | { 21 | "matchManagers": ["github-actions"], 22 | "groupName": "dependencies for github", 23 | "commitMessagePrefix": "chore(deps):", 24 | "pinDigest": true, 25 | }, 26 | { 27 | "groupName": "python-nonmajor", 28 | "matchLanguages": ["python"], 29 | "matchUpdateTypes": ["minor", "patch"], 30 | }, 31 | { 32 | "matchPackageNames": ["psycopg2-binary"], 33 | "allowedVersions": "!/2.9.10/" 34 | }, 35 | ], 36 | } 37 | -------------------------------------------------------------------------------- /.github/workflows/codeql.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: "CodeQL" 16 | 17 | on: 18 | push: 19 | branches: [ "main" ] 20 | pull_request: 21 | branches: [ "main" ] 22 | paths-ignore: 23 | - '**/*.md' 24 | - '**/*.txt' 25 | 26 | # Declare default permissions as read only. 27 | permissions: read-all 28 | 29 | jobs: 30 | analyze: 31 | name: Analyze 32 | runs-on: ubuntu-latest 33 | permissions: 34 | actions: read 35 | contents: read 36 | security-events: write 37 | 38 | strategy: 39 | fail-fast: false 40 | matrix: 41 | language: [ 'python' ] 42 | 43 | steps: 44 | - name: Checkout repository 45 | uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 46 | with: 47 | ref: ${{ github.event.pull_request.head.sha }} 48 | repository: ${{ github.event.pull_request.head.repo.full_name }} 49 | 50 | # Initializes the CodeQL tools for scanning. 51 | - name: Initialize CodeQL 52 | uses: github/codeql-action/init@181d5eefc20863364f96762470ba6f862bdef56b # v3.29.2 53 | with: 54 | languages: ${{ matrix.language }} 55 | 56 | # Autobuild attempts to build any compiled languages (C/C++, C#, Go, or Java). 57 | # If this step fails, then you should remove it and run the build manually 58 | - name: Autobuild 59 | uses: github/codeql-action/autobuild@181d5eefc20863364f96762470ba6f862bdef56b # v3.29.2 60 | 61 | - name: Perform CodeQL Analysis 62 | uses: github/codeql-action/analyze@181d5eefc20863364f96762470ba6f862bdef56b # v3.29.2 63 | with: 64 | category: "/language:${{matrix.language}}" 65 | -------------------------------------------------------------------------------- /.github/workflows/coverage.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: Code Coverage 16 | on: 17 | pull_request: 18 | branches: 19 | - main 20 | 21 | # Declare default permissions as read only. 22 | permissions: read-all 23 | 24 | jobs: 25 | coverage: 26 | runs-on: ubuntu-latest 27 | permissions: 28 | issues: write 29 | pull-requests: write 30 | steps: 31 | - name: Setup Python 32 | uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 33 | with: 34 | python-version: "3.13" 35 | 36 | - run: pip install nox coverage 37 | 38 | - name: Checkout base branch 39 | uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 40 | with: 41 | ref: ${{ github.base_ref }} 42 | 43 | - name: Calculate base code coverage 44 | run: | 45 | nox --sessions unit-3.13 46 | coverage report --show-missing 47 | export CUR_COVER=$(coverage report | awk '$1 == "TOTAL" {print $NF+0}') 48 | echo "CUR_COVER=$CUR_COVER" >> $GITHUB_ENV 49 | coverage erase 50 | 51 | - name: Checkout PR branch 52 | uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 53 | with: 54 | ref: ${{ github.event.pull_request.head.sha }} 55 | repository: ${{ github.event.pull_request.head.repo.full_name }} 56 | 57 | - name: Calculate PR code coverage 58 | run: | 59 | nox --sessions unit-3.13 60 | coverage report --show-missing 61 | export PR_COVER=$(coverage report | awk '$1 == "TOTAL" {print $NF+0}') 62 | echo "PR_COVER=$PR_COVER" >> $GITHUB_ENV 63 | coverage erase 64 | 65 | - name: Verify code coverage. If your reading this and the step has failed, please add tests to cover your changes. 66 | run: | 67 | echo "BASE BRANCH CODE COVERAGE is ${{ env.CUR_COVER }}%" 68 | echo "PULL REQUEST CODE COVERAGE is ${{ env.PR_COVER }}%" 69 | if [ "${{ env.PR_COVER }}" -lt "${{ env.CUR_COVER }}" ]; then 70 | exit 1; 71 | fi 72 | -------------------------------------------------------------------------------- /.github/workflows/labels.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: Sync labels 16 | on: 17 | push: 18 | branches: 19 | - main 20 | 21 | # Declare default permissions as read only. 22 | permissions: read-all 23 | 24 | jobs: 25 | build: 26 | runs-on: ubuntu-latest 27 | permissions: 28 | issues: write 29 | pull-requests: write 30 | steps: 31 | - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 32 | - uses: micnncim/action-label-syncer@3abd5ab72fda571e69fffd97bd4e0033dd5f495c # v1.3.0 33 | env: 34 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 35 | -------------------------------------------------------------------------------- /.github/workflows/lint.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: Lint 16 | on: 17 | pull_request: 18 | branches: 19 | - main 20 | 21 | # Declare default permissions as read only. 22 | permissions: read-all 23 | 24 | jobs: 25 | lint: 26 | name: Run lint 27 | runs-on: ubuntu-latest 28 | permissions: 29 | issues: write 30 | pull-requests: write 31 | steps: 32 | - name: Setup Python 33 | uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 34 | with: 35 | python-version: "3.13" 36 | 37 | - name: Install nox 38 | run: pip install nox 39 | 40 | - name: Checkout code 41 | uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 42 | with: 43 | ref: ${{ github.event.pull_request.head.sha }} 44 | repository: ${{ github.event.pull_request.head.repo.full_name }} 45 | 46 | - name: Run nox lint session 47 | run: nox -s lint 48 | -------------------------------------------------------------------------------- /.github/workflows/scorecard.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: OSSF Scorecard 16 | on: 17 | # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection 18 | branch_protection_rule: 19 | schedule: 20 | # weekly on Sunday 21 | - cron: '0 20 * * 0' 22 | push: 23 | branches: [ "main" ] 24 | 25 | # Declare default permissions as read only. 26 | permissions: read-all 27 | 28 | jobs: 29 | analysis: 30 | name: Scorecard analysis 31 | runs-on: ubuntu-latest 32 | permissions: 33 | # Needed to upload the results to code-scanning dashboard. 34 | security-events: write 35 | 36 | steps: 37 | - name: "Checkout code" 38 | uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 39 | with: 40 | persist-credentials: false 41 | 42 | - name: "Run analysis" 43 | uses: ossf/scorecard-action@05b42c624433fc40578a4040d5cf5e36ddca8cde # v2.4.2 44 | with: 45 | results_file: results.sarif 46 | results_format: sarif 47 | 48 | - name: Filter SARIF to skip false positives 49 | # filter out DangerousWorkflow alerts as they do not account for safe use of labels to trigger actions 50 | env: 51 | SCORECARD_SKIPPED_RULE_IDS: "DangerousWorkflowID" 52 | run: | 53 | SCORECARD_SKIPPED_RULE_IDS_JSON=$(echo $SCORECARD_SKIPPED_RULE_IDS | jq -cR 'split(",")') 54 | # Trim the SARIF file to remove false positive detections 55 | cat results.sarif | jq '.runs[].results |= map(select(.ruleId as $id | '$SCORECARD_SKIPPED_RULE_IDS_JSON' | all($id != .)))' > resultsFiltered.sarif 56 | 57 | # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF 58 | # format to the repository Actions tab. 59 | - name: "Upload artifact" 60 | uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 61 | with: 62 | name: SARIF file 63 | path: results.sarif 64 | retention-days: 5 65 | 66 | # Upload the results to GitHub's code scanning dashboard. 67 | - name: "Upload to code-scanning" 68 | uses: github/codeql-action/upload-sarif@181d5eefc20863364f96762470ba6f862bdef56b # v3.29.2 69 | with: 70 | sarif_file: resultsFiltered.sarif 71 | -------------------------------------------------------------------------------- /.github/workflows/tests.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: tests 16 | on: 17 | pull_request: 18 | branches: 19 | - main 20 | push: 21 | branches: 22 | - main 23 | schedule: 24 | - cron: '0 2 * * *' 25 | 26 | # Declare default permissions as read only. 27 | permissions: read-all 28 | 29 | jobs: 30 | unit: 31 | name: unit tests 32 | runs-on: ${{ matrix.os }} 33 | permissions: 34 | contents: read 35 | id-token: write 36 | issues: write 37 | pull-requests: write 38 | strategy: 39 | matrix: 40 | os: [macos-latest, windows-latest, ubuntu-latest] 41 | python-version: ["3.9", "3.13"] 42 | fail-fast: false 43 | steps: 44 | - name: Checkout code 45 | uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 46 | with: 47 | ref: ${{ github.event.pull_request.head.sha }} 48 | repository: ${{ github.event.pull_request.head.repo.full_name }} 49 | 50 | - name: Setup Python ${{ matrix.python-version }} 51 | uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 52 | with: 53 | python-version: ${{ matrix.python-version }} 54 | 55 | - name: Install nox 56 | run: pip install nox 57 | 58 | - id: 'auth' 59 | name: Authenticate to Google Cloud 60 | # only needed for Flakybot on periodic (schedule) and continuous (push) events 61 | if: ${{ github.event_name == 'schedule' || github.event_name == 'push' }} 62 | uses: google-github-actions/auth@ba79af03959ebeac9769e648f473a284504d9193 # v2.1.10 63 | with: 64 | workload_identity_provider: ${{ vars.PROVIDER_NAME }} 65 | service_account: ${{ vars.SERVICE_ACCOUNT }} 66 | access_token_lifetime: 600s 67 | 68 | - name: Run tests 69 | run: nox -s unit-${{ matrix.python-version }} 70 | 71 | - name: FlakyBot (Linux) 72 | # only run flakybot on periodic (schedule) and continuous (push) events 73 | if: ${{ (github.event_name == 'schedule' || github.event_name == 'push') && runner.os == 'Linux' && always() }} 74 | run: | 75 | curl https://github.com/googleapis/repo-automation-bots/releases/download/flakybot-1.1.0/flakybot -o flakybot -s -L 76 | chmod +x ./flakybot 77 | ./flakybot --repo ${{github.repository}} --commit_hash ${{github.sha}} --build_url https://github.com/${{github.repository}}/actions/runs/${{github.run_id}} 78 | - name: FlakyBot (Windows) 79 | # only run flakybot on periodic (schedule) and continuous (push) events 80 | if: ${{ (github.event_name == 'schedule' || github.event_name == 'push') && runner.os == 'Windows' && always() }} 81 | run: | 82 | curl https://github.com/googleapis/repo-automation-bots/releases/download/flakybot-1.1.0/flakybot.exe -o flakybot.exe -s -L 83 | ./flakybot.exe --repo ${{github.repository}} --commit_hash ${{github.sha}} --build_url https://github.com/${{github.repository}}/actions/runs/${{github.run_id}} 84 | - name: FlakyBot (macOS) 85 | # only run flakybot on periodic (schedule) and continuous (push) events 86 | if: ${{ (github.event_name == 'schedule' || github.event_name == 'push') && runner.os == 'macOS' && always() }} 87 | run: | 88 | curl https://github.com/googleapis/repo-automation-bots/releases/download/flakybot-1.1.0/flakybot-darwin-amd64 -o flakybot -s -L 89 | chmod +x ./flakybot 90 | ./flakybot --repo ${{github.repository}} --commit_hash ${{github.sha}} --build_url https://github.com/${{github.repository}}/actions/runs/${{github.run_id}} 91 | 92 | integration: 93 | name: integration tests 94 | runs-on: [self-hosted, linux, x64] 95 | # run integration tests on all builds except pull requests from forks or 96 | # dependabot 97 | if: | 98 | github.event_name != 'pull_request' || 99 | (github.event.pull_request.head.repo.full_name == github.repository && github.actor != 'dependabot[bot]') 100 | strategy: 101 | matrix: 102 | python-version: ["3.9", "3.13"] 103 | fail-fast: false 104 | permissions: 105 | contents: read 106 | id-token: write 107 | issues: write 108 | pull-requests: write 109 | steps: 110 | - name: Checkout code 111 | uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 112 | with: 113 | ref: ${{ github.event.pull_request.head.sha }} 114 | repository: ${{ github.event.pull_request.head.repo.full_name }} 115 | 116 | - name: Setup Python ${{ matrix.python-version }} 117 | uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 118 | with: 119 | python-version: ${{ matrix.python-version }} 120 | 121 | - name: Install nox 122 | run: pip install nox 123 | 124 | - id: 'auth' 125 | name: 'Authenticate to Google Cloud' 126 | uses: google-github-actions/auth@ba79af03959ebeac9769e648f473a284504d9193 # v2.1.10 127 | with: 128 | workload_identity_provider: ${{ vars.PROVIDER_NAME }} 129 | service_account: ${{ vars.SERVICE_ACCOUNT }} 130 | access_token_lifetime: 600s 131 | 132 | - id: 'secrets' 133 | name: Get secrets 134 | uses: google-github-actions/get-secretmanager-secrets@a8440875e1c2892062aef9061228d4f1af8f919b # v2.2.3 135 | with: 136 | secrets: |- 137 | ALLOYDB_INSTANCE_URI:${{ vars.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_INSTANCE_URI 138 | ALLOYDB_CLUSTER_PASS:${{ vars.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_CLUSTER_PASS 139 | ALLOYDB_IAM_USER:${{ vars.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_PYTHON_IAM_USER 140 | ALLOYDB_INSTANCE_IP:${{ vars.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_INSTANCE_IP 141 | ALLOYDB_PSC_INSTANCE_URI:${{ vars.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_PSC_INSTANCE_URI 142 | 143 | - name: Run tests 144 | env: 145 | ALLOYDB_DB: 'postgres' 146 | ALLOYDB_USER: 'postgres' 147 | ALLOYDB_PASS: '${{ steps.secrets.outputs.ALLOYDB_CLUSTER_PASS }}' 148 | ALLOYDB_IAM_USER: '${{ steps.secrets.outputs.ALLOYDB_IAM_USER }}' 149 | ALLOYDB_INSTANCE_IP: '${{ steps.secrets.outputs.ALLOYDB_INSTANCE_IP }}' 150 | ALLOYDB_INSTANCE_URI: '${{ steps.secrets.outputs.ALLOYDB_INSTANCE_URI }}' 151 | ALLOYDB_PSC_INSTANCE_URI: '${{ steps.secrets.outputs.ALLOYDB_PSC_INSTANCE_URI }}' 152 | run: nox -s system-${{ matrix.python-version }} 153 | 154 | - name: FlakyBot (Linux) 155 | # only run flakybot on periodic (schedule) and continuous (push) events 156 | if: ${{ (github.event_name == 'schedule' || github.event_name == 'push') && always() }} 157 | run: | 158 | curl https://github.com/googleapis/repo-automation-bots/releases/download/flakybot-1.1.0/flakybot -o flakybot -s -L 159 | chmod +x ./flakybot 160 | ./flakybot --repo ${{github.repository}} --commit_hash ${{github.sha}} --build_url https://github.com/${{github.repository}}/actions/runs/${{github.run_id}} 161 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.py[cod] 2 | *.sw[op] 3 | 4 | # C extensions 5 | *.so 6 | 7 | # Packages 8 | *.egg 9 | *.egg-info 10 | dist 11 | build 12 | eggs 13 | .eggs 14 | parts 15 | bin 16 | var 17 | sdist 18 | develop-eggs 19 | .installed.cfg 20 | lib 21 | lib64 22 | __pycache__ 23 | 24 | # Installer logs 25 | pip-log.txt 26 | 27 | # Unit test / coverage reports 28 | .coverage 29 | .nox 30 | .cache 31 | .pytest_cache 32 | 33 | 34 | # Mac 35 | .DS_Store 36 | 37 | # JetBrains 38 | .idea 39 | 40 | # VS Code 41 | .vscode 42 | 43 | # emacs 44 | *~ 45 | 46 | # Built documentation 47 | docs/_build 48 | bigquery/docs/generated 49 | docs.metadata 50 | 51 | # Virtual environment 52 | env/ 53 | 54 | # Test logs 55 | coverage.xml 56 | *sponge_log.xml 57 | 58 | # System test environment variables. 59 | system_tests/local_test_setup 60 | 61 | # Make sure a generated file isn't accidentally committed. 62 | pylintrc 63 | pylintrc.test 64 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to making participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, gender identity and expression, level of 9 | experience, education, socio-economic status, nationality, personal appearance, 10 | race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or reject 41 | comments, commits, code, wiki edits, issues, and other contributions that are 42 | not aligned to this Code of Conduct, or to ban temporarily or permanently any 43 | contributor for other behaviors that they deem inappropriate, threatening, 44 | offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies both within project spaces and in public spaces 49 | when an individual is representing the project or its community. Examples of 50 | representing a project or community include using an official project e-mail 51 | address, posting via an official social media account, or acting as an appointed 52 | representative at an online or offline event. Representation of a project may be 53 | further defined and clarified by project maintainers. 54 | 55 | This Code of Conduct also applies outside the project spaces when the Project 56 | Steward has a reasonable belief that an individual's behavior may have a 57 | negative impact on the project or its community. 58 | 59 | ## Conflict Resolution 60 | 61 | We do not believe that all conflict is bad; healthy debate and disagreement 62 | often yield positive results. However, it is never okay to be disrespectful or 63 | to engage in behavior that violates the project’s code of conduct. 64 | 65 | If you see someone violating the code of conduct, you are encouraged to address 66 | the behavior directly with those involved. Many issues can be resolved quickly 67 | and easily, and this gives people more control over the outcome of their 68 | dispute. If you are unable to resolve the matter for any reason, or if the 69 | behavior is threatening or harassing, report it. We are dedicated to providing 70 | an environment where participants feel welcome and safe. 71 | 72 | Reports should be directed to *googleapis-stewards@google.com*, the 73 | Project Steward(s) for *Google Cloud Client Libraries*. It is the Project Steward’s duty to 74 | receive and address reported violations of the code of conduct. They will then 75 | work with a committee consisting of representatives from the Open Source 76 | Programs Office and the Google Open Source Strategy team. If for any reason you 77 | are uncomfortable reaching out to the Project Steward, please email 78 | opensource@google.com. 79 | 80 | We will investigate every complaint, but you may not receive a direct response. 81 | We will use our discretion in determining when and how to follow up on reported 82 | incidents, which may range from not taking action to permanent expulsion from 83 | the project and project-sponsored spaces. We will notify the accused of the 84 | report and provide them an opportunity to discuss it before any action is taken. 85 | The identity of the reporter will be omitted from the details of the report 86 | supplied to the accused. In potentially harmful situations, such as ongoing 87 | harassment or threats to anyone's safety, we may take action without notice. 88 | 89 | ## Attribution 90 | 91 | This Code of Conduct is adapted from the Contributor Covenant, version 1.4, 92 | available at 93 | https://www.contributor-covenant.org/version/1/4/code-of-conduct.htmls 94 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # How to Contribute 2 | 3 | We'd love to accept your patches and contributions to this project. There are 4 | just a few small guidelines you need to follow. 5 | 6 | ## Contributor License Agreement 7 | 8 | Contributions to this project must be accompanied by a Contributor License 9 | Agreement. You (or your employer) retain the copyright to your contribution; 10 | this simply gives us permission to use and redistribute your contributions as 11 | part of the project. Head over to to see 12 | your current agreements on file or to sign a new one. 13 | 14 | You generally only need to submit a CLA once, so if you've already submitted one 15 | (even if it was for a different project), you probably don't need to do it 16 | again. 17 | 18 | ## Code Reviews 19 | 20 | All submissions, including submissions by project members, require review. We 21 | use GitHub pull requests for this purpose. Consult 22 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more 23 | information on using pull requests. 24 | 25 | ## Community Guidelines 26 | 27 | This project follows [Google's Open Source Community 28 | Guidelines](https://opensource.google/conduct/). 29 | 30 | ## Testing 31 | 32 | NOTE: Be sure to run the following commands in the same VPC as the AlloyDB instance 33 | 34 | 1. Setup local environment to work with virtualenv and nox if you haven't already, e.g. 35 | ``` 36 | python3 -m venv venv 37 | source ./venv/bin/activate 38 | pip install nox 39 | ``` 40 | 1. Set the environment variables. You can see an example of the environment variables needed by running `cat .envrc.example` 41 | 1. Run `gcloud auth application-default login` 42 | 1. Command to run the unit tests: `nox -s unit-` 43 | 1. Command to run the integration tests: `nox -s system-` 44 | -------------------------------------------------------------------------------- /SECURITY.md: -------------------------------------------------------------------------------- 1 | # Security Policy 2 | 3 | To report a security issue, please use [g.co/vulnz](https://g.co/vulnz). 4 | 5 | The Google Security Team will respond within 5 working days of your report on [g.co/vulnz](https://g.co/vulnz). 6 | 7 | We use [g.co/vulnz](https://g.co/vulnz) for our intake, and do coordination and disclosure here using GitHub Security Advisory to privately discuss and fix the issue. 8 | -------------------------------------------------------------------------------- /docs/images/alloydb-python-connector.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GoogleCloudPlatform/alloydb-python-connector/86d38bbff1c409cae84696a81bb009a71a3e2a9d/docs/images/alloydb-python-connector.png -------------------------------------------------------------------------------- /google/api/field_behavior_pb2.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # type: ignore 16 | # -*- coding: utf-8 -*- 17 | # Generated by the protocol buffer compiler. DO NOT EDIT! 18 | # source: google/api/field_behavior.proto 19 | # isort: skip_file 20 | """Generated protocol buffer code.""" 21 | from google.protobuf import descriptor as _descriptor 22 | from google.protobuf import descriptor_pool as _descriptor_pool 23 | from google.protobuf import symbol_database as _symbol_database 24 | from google.protobuf.internal import builder as _builder 25 | 26 | # @@protoc_insertion_point(imports) 27 | 28 | _sym_db = _symbol_database.Default() 29 | 30 | 31 | from google.protobuf import descriptor_pb2 as google_dot_protobuf_dot_descriptor__pb2 32 | 33 | 34 | DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( 35 | b"\n\x1fgoogle/api/field_behavior.proto\x12\ngoogle.api\x1a google/protobuf/descriptor.proto*\xb6\x01\n\rFieldBehavior\x12\x1e\n\x1a\x46IELD_BEHAVIOR_UNSPECIFIED\x10\x00\x12\x0c\n\x08OPTIONAL\x10\x01\x12\x0c\n\x08REQUIRED\x10\x02\x12\x0f\n\x0bOUTPUT_ONLY\x10\x03\x12\x0e\n\nINPUT_ONLY\x10\x04\x12\r\n\tIMMUTABLE\x10\x05\x12\x12\n\x0eUNORDERED_LIST\x10\x06\x12\x15\n\x11NON_EMPTY_DEFAULT\x10\x07\x12\x0e\n\nIDENTIFIER\x10\x08:Q\n\x0e\x66ield_behavior\x12\x1d.google.protobuf.FieldOptions\x18\x9c\x08 \x03(\x0e\x32\x19.google.api.FieldBehaviorBp\n\x0e\x63om.google.apiB\x12\x46ieldBehaviorProtoP\x01ZAgoogle.golang.org/genproto/googleapis/api/annotations;annotations\xa2\x02\x04GAPIb\x06proto3" 36 | ) 37 | 38 | _globals = globals() 39 | _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) 40 | _builder.BuildTopDescriptorsAndMessages( 41 | DESCRIPTOR, "google.api.field_behavior_pb2", _globals 42 | ) 43 | if _descriptor._USE_C_DESCRIPTORS == False: 44 | google_dot_protobuf_dot_descriptor__pb2.FieldOptions.RegisterExtension( 45 | field_behavior 46 | ) 47 | 48 | DESCRIPTOR._options = None 49 | DESCRIPTOR._serialized_options = b"\n\016com.google.apiB\022FieldBehaviorProtoP\001ZAgoogle.golang.org/genproto/googleapis/api/annotations;annotations\242\002\004GAPI" 50 | _globals["_FIELDBEHAVIOR"]._serialized_start = 82 51 | _globals["_FIELDBEHAVIOR"]._serialized_end = 264 52 | # @@protoc_insertion_point(module_scope) 53 | -------------------------------------------------------------------------------- /google/cloud/alloydb/connector/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from google.cloud.alloydbconnector import __version__ 15 | from google.cloud.alloydbconnector import AsyncConnector 16 | from google.cloud.alloydbconnector import Connector 17 | from google.cloud.alloydbconnector import IPTypes 18 | from google.cloud.alloydbconnector import RefreshStrategy 19 | 20 | __all__ = [ 21 | "__version__", 22 | "Connector", 23 | "AsyncConnector", 24 | "IPTypes", 25 | "RefreshStrategy", 26 | ] 27 | -------------------------------------------------------------------------------- /google/cloud/alloydb_connectors_v1/proto/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /google/cloud/alloydb_connectors_v1/proto/resources_pb2.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # -*- coding: utf-8 -*- 16 | # Generated by the protocol buffer compiler. DO NOT EDIT! 17 | # source: google/cloud/alloydb_connectors_v1/proto/resources.proto 18 | # isort: skip_file 19 | """Generated protocol buffer code.""" 20 | from google.protobuf import descriptor as _descriptor 21 | from google.protobuf import descriptor_pool as _descriptor_pool 22 | from google.protobuf import symbol_database as _symbol_database 23 | from google.protobuf.internal import builder as _builder 24 | 25 | # @@protoc_insertion_point(imports) 26 | 27 | _sym_db = _symbol_database.Default() 28 | 29 | 30 | from google.api import field_behavior_pb2 as google_dot_api_dot_field__behavior__pb2 31 | 32 | 33 | DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( 34 | b'\n8google/cloud/alloydb_connectors_v1/proto/resources.proto\x12"google.cloud.alloydb.connectors.v1\x1a\x1fgoogle/api/field_behavior.proto"\xe6\x01\n\x17MetadataExchangeRequest\x12\x18\n\nuser_agent\x18\x01 \x01(\tB\x04\xe2\x41\x01\x01\x12W\n\tauth_type\x18\x02 \x01(\x0e\x32\x44.google.cloud.alloydb.connectors.v1.MetadataExchangeRequest.AuthType\x12\x14\n\x0coauth2_token\x18\x03 \x01(\t"B\n\x08\x41uthType\x12\x19\n\x15\x41UTH_TYPE_UNSPECIFIED\x10\x00\x12\r\n\tDB_NATIVE\x10\x01\x12\x0c\n\x08\x41UTO_IAM\x10\x02"\xd3\x01\n\x18MetadataExchangeResponse\x12`\n\rresponse_code\x18\x01 \x01(\x0e\x32I.google.cloud.alloydb.connectors.v1.MetadataExchangeResponse.ResponseCode\x12\x13\n\x05\x65rror\x18\x02 \x01(\tB\x04\xe2\x41\x01\x01"@\n\x0cResponseCode\x12\x1d\n\x19RESPONSE_CODE_UNSPECIFIED\x10\x00\x12\x06\n\x02OK\x10\x01\x12\t\n\x05\x45RROR\x10\x02\x42\xf5\x01\n&com.google.cloud.alloydb.connectors.v1B\x0eResourcesProtoP\x01ZFcloud.google.com/go/alloydb/connectors/apiv1/connectorspb;connectorspb\xaa\x02"Google.Cloud.AlloyDb.Connectors.V1\xca\x02"Google\\Cloud\\AlloyDb\\Connectors\\V1\xea\x02&Google::Cloud::AlloyDb::Connectors::V1b\x06proto3' 35 | ) 36 | 37 | _globals = globals() 38 | _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) 39 | _builder.BuildTopDescriptorsAndMessages( 40 | DESCRIPTOR, "google.cloud.alloydb_connectors_v1.proto.resources_pb2", _globals 41 | ) 42 | if _descriptor._USE_C_DESCRIPTORS == False: 43 | DESCRIPTOR._options = None 44 | DESCRIPTOR._serialized_options = b'\n&com.google.cloud.alloydb.connectors.v1B\016ResourcesProtoP\001ZFcloud.google.com/go/alloydb/connectors/apiv1/connectorspb;connectorspb\252\002"Google.Cloud.AlloyDb.Connectors.V1\312\002"Google\\Cloud\\AlloyDb\\Connectors\\V1\352\002&Google::Cloud::AlloyDb::Connectors::V1' 45 | _METADATAEXCHANGEREQUEST.fields_by_name["user_agent"]._options = None 46 | _METADATAEXCHANGEREQUEST.fields_by_name[ 47 | "user_agent" 48 | ]._serialized_options = b"\342A\001\001" 49 | _METADATAEXCHANGERESPONSE.fields_by_name["error"]._options = None 50 | _METADATAEXCHANGERESPONSE.fields_by_name[ 51 | "error" 52 | ]._serialized_options = b"\342A\001\001" 53 | _globals["_METADATAEXCHANGEREQUEST"]._serialized_start = 130 54 | _globals["_METADATAEXCHANGEREQUEST"]._serialized_end = 360 55 | _globals["_METADATAEXCHANGEREQUEST_AUTHTYPE"]._serialized_start = 294 56 | _globals["_METADATAEXCHANGEREQUEST_AUTHTYPE"]._serialized_end = 360 57 | _globals["_METADATAEXCHANGERESPONSE"]._serialized_start = 363 58 | _globals["_METADATAEXCHANGERESPONSE"]._serialized_end = 574 59 | _globals["_METADATAEXCHANGERESPONSE_RESPONSECODE"]._serialized_start = 510 60 | _globals["_METADATAEXCHANGERESPONSE_RESPONSECODE"]._serialized_end = 574 61 | # @@protoc_insertion_point(module_scope) 62 | -------------------------------------------------------------------------------- /google/cloud/alloydb_connectors_v1/proto/resources_pb2.pyi: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from typing import ClassVar as _ClassVar 16 | from typing import Optional as _Optional 17 | from typing import Union as _Union 18 | 19 | from google.protobuf import descriptor as _descriptor 20 | from google.protobuf import message as _message 21 | from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper 22 | 23 | from google.api import field_behavior_pb2 as _field_behavior_pb2 24 | 25 | DESCRIPTOR: _descriptor.FileDescriptor 26 | 27 | class MetadataExchangeRequest(_message.Message): 28 | __slots__ = ["auth_type", "oauth2_token", "user_agent"] 29 | 30 | class AuthType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): 31 | __slots__ = [] # type: ignore 32 | AUTH_TYPE_FIELD_NUMBER: _ClassVar[int] 33 | AUTH_TYPE_UNSPECIFIED: MetadataExchangeRequest.AuthType 34 | AUTO_IAM: MetadataExchangeRequest.AuthType 35 | DB_NATIVE: MetadataExchangeRequest.AuthType 36 | OAUTH2_TOKEN_FIELD_NUMBER: _ClassVar[int] 37 | USER_AGENT_FIELD_NUMBER: _ClassVar[int] 38 | auth_type: MetadataExchangeRequest.AuthType 39 | oauth2_token: str 40 | user_agent: str 41 | def __init__( 42 | self, 43 | user_agent: _Optional[str] = ..., 44 | auth_type: _Optional[_Union[MetadataExchangeRequest.AuthType, str]] = ..., 45 | oauth2_token: _Optional[str] = ..., 46 | ) -> None: ... 47 | 48 | class MetadataExchangeResponse(_message.Message): 49 | __slots__ = ["error", "response_code"] 50 | 51 | class ResponseCode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): 52 | __slots__ = [] # type: ignore 53 | ERROR: MetadataExchangeResponse.ResponseCode 54 | ERROR_FIELD_NUMBER: _ClassVar[int] 55 | OK: MetadataExchangeResponse.ResponseCode 56 | RESPONSE_CODE_FIELD_NUMBER: _ClassVar[int] 57 | RESPONSE_CODE_UNSPECIFIED: MetadataExchangeResponse.ResponseCode 58 | error: str 59 | response_code: MetadataExchangeResponse.ResponseCode 60 | def __init__( 61 | self, 62 | response_code: _Optional[ 63 | _Union[MetadataExchangeResponse.ResponseCode, str] 64 | ] = ..., 65 | error: _Optional[str] = ..., 66 | ) -> None: ... 67 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from google.cloud.alloydbconnector.async_connector import AsyncConnector 15 | from google.cloud.alloydbconnector.connector import Connector 16 | from google.cloud.alloydbconnector.enums import IPTypes 17 | from google.cloud.alloydbconnector.enums import RefreshStrategy 18 | from google.cloud.alloydbconnector.version import __version__ 19 | 20 | __all__ = [ 21 | "__version__", 22 | "Connector", 23 | "AsyncConnector", 24 | "IPTypes", 25 | "RefreshStrategy", 26 | ] 27 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/async_connector.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from __future__ import annotations 16 | 17 | import asyncio 18 | import logging 19 | from types import TracebackType 20 | from typing import Any, Optional, TYPE_CHECKING 21 | 22 | import google.auth 23 | from google.auth.credentials import with_scopes_if_required 24 | import google.auth.transport.requests 25 | 26 | import google.cloud.alloydbconnector.asyncpg as asyncpg 27 | from google.cloud.alloydbconnector.client import AlloyDBClient 28 | from google.cloud.alloydbconnector.enums import IPTypes 29 | from google.cloud.alloydbconnector.enums import RefreshStrategy 30 | from google.cloud.alloydbconnector.exceptions import ClosedConnectorError 31 | from google.cloud.alloydbconnector.instance import RefreshAheadCache 32 | from google.cloud.alloydbconnector.lazy import LazyRefreshCache 33 | from google.cloud.alloydbconnector.types import CacheTypes 34 | from google.cloud.alloydbconnector.utils import generate_keys 35 | from google.cloud.alloydbconnector.utils import strip_http_prefix 36 | 37 | if TYPE_CHECKING: 38 | from google.auth.credentials import Credentials 39 | 40 | logger = logging.getLogger(name=__name__) 41 | 42 | 43 | class AsyncConnector: 44 | """A class to configure and create connections to Cloud SQL instances 45 | asynchronously. 46 | 47 | Args: 48 | credentials (google.auth.credentials.Credentials): 49 | A credentials object created from the google-auth Python library. 50 | If not specified, Application Default Credentials are used. 51 | quota_project (str): The Project ID for an existing Google Cloud 52 | project. The project specified is used for quota and 53 | billing purposes. 54 | Defaults to None, picking up project from environment. 55 | alloydb_api_endpoint (str): Base URL to use when calling 56 | the AlloyDB API endpoint. Defaults to "alloydb.googleapis.com". 57 | enable_iam_auth (bool): Enables automatic IAM database authentication. 58 | ip_type (str | IPTypes): Default IP type for all AlloyDB connections. 59 | Defaults to IPTypes.PRIVATE ("PRIVATE") for private IP connections. 60 | refresh_strategy (str | RefreshStrategy): The default refresh strategy 61 | used to refresh SSL/TLS cert and instance metadata. Can be one 62 | of the following: RefreshStrategy.LAZY ("LAZY") or 63 | RefreshStrategy.BACKGROUND ("BACKGROUND"). 64 | Default: RefreshStrategy.BACKGROUND 65 | """ 66 | 67 | def __init__( 68 | self, 69 | credentials: Optional[Credentials] = None, 70 | quota_project: Optional[str] = None, 71 | alloydb_api_endpoint: str = "alloydb.googleapis.com", 72 | enable_iam_auth: bool = False, 73 | ip_type: str | IPTypes = IPTypes.PRIVATE, 74 | user_agent: Optional[str] = None, 75 | refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND, 76 | ) -> None: 77 | self._cache: dict[str, CacheTypes] = {} 78 | # initialize default params 79 | self._quota_project = quota_project 80 | self._alloydb_api_endpoint = strip_http_prefix(alloydb_api_endpoint) 81 | self._enable_iam_auth = enable_iam_auth 82 | # if ip_type is str, convert to IPTypes enum 83 | if isinstance(ip_type, str): 84 | ip_type = IPTypes(ip_type.upper()) 85 | self._ip_type = ip_type 86 | # if refresh_strategy is str, convert to RefreshStrategy enum 87 | if isinstance(refresh_strategy, str): 88 | refresh_strategy = RefreshStrategy(refresh_strategy.upper()) 89 | self._refresh_strategy = refresh_strategy 90 | self._user_agent = user_agent 91 | # initialize credentials 92 | scopes = ["https://www.googleapis.com/auth/cloud-platform"] 93 | if credentials: 94 | self._credentials = with_scopes_if_required(credentials, scopes=scopes) 95 | # otherwise use application default credentials 96 | else: 97 | self._credentials, _ = google.auth.default(scopes=scopes) 98 | 99 | # check if AsyncConnector is being initialized with event loop running 100 | # Otherwise we will lazy init keys 101 | try: 102 | self._keys: Optional[asyncio.Task] = asyncio.create_task(generate_keys()) 103 | except RuntimeError: 104 | self._keys = None 105 | self._client: Optional[AlloyDBClient] = None 106 | self._closed = False 107 | 108 | async def connect( 109 | self, 110 | instance_uri: str, 111 | driver: str, 112 | **kwargs: Any, 113 | ) -> Any: 114 | """ 115 | Asynchronously prepares and returns a database connection object. 116 | 117 | Starts tasks to refresh the certificates and get 118 | AlloyDB instance IP address. Creates a secure TLS connection 119 | to establish connection to AlloyDB instance. 120 | 121 | Args: 122 | instance_uri (str): The instance URI of the AlloyDB instance. 123 | ex. projects//locations//clusters//instances/ 124 | driver (str): A string representing the database driver to connect 125 | with. Supported drivers are asyncpg. 126 | **kwargs: Pass in any database driver-specific arguments needed 127 | to fine tune connection. 128 | 129 | Returns: 130 | connection: A DBAPI connection to the specified AlloyDB instance. 131 | """ 132 | if self._closed: 133 | raise ClosedConnectorError( 134 | "Connection attempt failed because the connector has already been closed." 135 | ) 136 | if self._keys is None: 137 | self._keys = asyncio.create_task(generate_keys()) 138 | if self._client is None: 139 | # lazy init client as it has to be initialized in async context 140 | self._client = AlloyDBClient( 141 | self._alloydb_api_endpoint, 142 | self._quota_project, 143 | self._credentials, 144 | user_agent=self._user_agent, 145 | driver=driver, 146 | ) 147 | 148 | enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth) 149 | 150 | # use existing connection info if possible 151 | if instance_uri in self._cache: 152 | cache = self._cache[instance_uri] 153 | else: 154 | if self._refresh_strategy == RefreshStrategy.LAZY: 155 | logger.debug( 156 | f"['{instance_uri}']: Refresh strategy is set to lazy refresh" 157 | ) 158 | cache = LazyRefreshCache(instance_uri, self._client, self._keys) 159 | else: 160 | logger.debug( 161 | f"['{instance_uri}']: Refresh strategy is set to background" 162 | " refresh" 163 | ) 164 | cache = RefreshAheadCache(instance_uri, self._client, self._keys) 165 | self._cache[instance_uri] = cache 166 | logger.debug(f"['{instance_uri}']: Connection info added to cache") 167 | 168 | connect_func = { 169 | "asyncpg": asyncpg.connect, 170 | } 171 | # only accept supported database drivers 172 | try: 173 | connector = connect_func[driver] 174 | except KeyError: 175 | raise ValueError(f"Driver '{driver}' is not a supported database driver.") 176 | 177 | # Host and ssl options come from the certificates and instance IP 178 | # address so we don't want the user to specify them. 179 | kwargs.pop("host", None) 180 | kwargs.pop("ssl", None) 181 | kwargs.pop("port", None) 182 | 183 | # get connection info for AlloyDB instance 184 | ip_type: str | IPTypes = kwargs.pop("ip_type", self._ip_type) 185 | # if ip_type is str, convert to IPTypes enum 186 | if isinstance(ip_type, str): 187 | ip_type = IPTypes(ip_type.upper()) 188 | try: 189 | conn_info = await cache.connect_info() 190 | ip_address = conn_info.get_preferred_ip(ip_type) 191 | except Exception: 192 | # with an error from AlloyDB API call or IP type, invalidate the 193 | # cache and re-raise the error 194 | await self._remove_cached(instance_uri) 195 | raise 196 | logger.debug(f"['{instance_uri}']: Connecting to {ip_address}:5433") 197 | 198 | # callable to be used for auto IAM authn 199 | def get_authentication_token() -> str: 200 | """Get OAuth2 access token to be used for IAM database authentication""" 201 | # refresh credentials if expired 202 | if not self._credentials.valid: 203 | request = google.auth.transport.requests.Request() 204 | self._credentials.refresh(request) 205 | return self._credentials.token 206 | 207 | # if enable_iam_auth is set, use auth token as database password 208 | if enable_iam_auth: 209 | kwargs["password"] = get_authentication_token 210 | try: 211 | return await connector( 212 | ip_address, await conn_info.create_ssl_context(), **kwargs 213 | ) 214 | except Exception: 215 | # we attempt a force refresh, then throw the error 216 | await cache.force_refresh() 217 | raise 218 | 219 | async def _remove_cached(self, instance_uri: str) -> None: 220 | """Stops all background refreshes and deletes the connection 221 | info cache from the map of caches. 222 | """ 223 | logger.debug(f"['{instance_uri}']: Removing connection info from cache") 224 | # remove cache from stored caches and close it 225 | cache = self._cache.pop(instance_uri) 226 | await cache.close() 227 | 228 | async def __aenter__(self) -> Any: 229 | """Enter async context manager by returning Connector object""" 230 | return self 231 | 232 | async def __aexit__( 233 | self, 234 | exc_type: Optional[type[BaseException]], 235 | exc_val: Optional[BaseException], 236 | exc_tb: Optional[TracebackType], 237 | ) -> None: 238 | """Exit async context manager by closing Connector""" 239 | await self.close() 240 | 241 | async def close(self) -> None: 242 | """Helper function to cancel RefreshAheadCaches' tasks 243 | and close client.""" 244 | await asyncio.gather(*[cache.close() for cache in self._cache.values()]) 245 | self._closed = True 246 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/asyncpg.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import ssl 16 | from typing import Any, TYPE_CHECKING 17 | 18 | SERVER_PROXY_PORT = 5433 19 | 20 | if TYPE_CHECKING: 21 | import asyncpg 22 | 23 | 24 | async def connect( 25 | ip_address: str, ctx: ssl.SSLContext, **kwargs: Any 26 | ) -> "asyncpg.Connection": 27 | """Helper function to create an asyncpg DB-API connection object. 28 | 29 | :type ip_address: str 30 | :param ip_address: A string containing an IP address for the AlloyDB 31 | instance. 32 | 33 | :type ctx: ssl.SSLContext 34 | :param ctx: An SSLContext object created from the AlloyDB server CA 35 | cert and ephemeral cert. 36 | 37 | :type kwargs: Any 38 | :param kwargs: Keyword arguments for establishing asyncpg connection 39 | object to AlloyDB instance. 40 | 41 | :rtype: asyncpg.Connection 42 | :returns: An asyncpg.Connection object to an AlloyDB instance. 43 | """ 44 | try: 45 | import asyncpg 46 | except ImportError: 47 | raise ImportError( 48 | 'Unable to import module "asyncpg." Please install and try again.' 49 | ) 50 | user = kwargs.pop("user") 51 | db = kwargs.pop("db") 52 | passwd = kwargs.pop("password") 53 | 54 | return await asyncpg.connect( 55 | user=user, 56 | database=db, 57 | password=passwd, 58 | host=ip_address, 59 | port=SERVER_PROXY_PORT, 60 | ssl=ctx, 61 | direct_tls=True, 62 | **kwargs, 63 | ) 64 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/client.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from __future__ import annotations 16 | 17 | import asyncio 18 | import logging 19 | from typing import Optional, TYPE_CHECKING 20 | 21 | from cryptography import x509 22 | from google.api_core.client_options import ClientOptions 23 | from google.api_core.gapic_v1.client_info import ClientInfo 24 | from google.auth.credentials import TokenState 25 | from google.auth.transport import requests 26 | import google.cloud.alloydb_v1beta as v1beta 27 | from google.protobuf import duration_pb2 28 | 29 | from google.cloud.alloydbconnector.connection_info import ConnectionInfo 30 | from google.cloud.alloydbconnector.version import __version__ as version 31 | 32 | if TYPE_CHECKING: 33 | from google.auth.credentials import Credentials 34 | 35 | USER_AGENT: str = f"alloydb-python-connector/{version}" 36 | API_VERSION: str = "v1beta" 37 | 38 | logger = logging.getLogger(name=__name__) 39 | 40 | 41 | def _format_user_agent( 42 | driver: Optional[str], 43 | custom_user_agent: Optional[str], 44 | ) -> str: 45 | """ 46 | Appends user-defined user agents to the base default agent. 47 | """ 48 | agent = f"{USER_AGENT}+{driver}" if driver else USER_AGENT 49 | if custom_user_agent and isinstance(custom_user_agent, str): 50 | agent = f"{agent} {custom_user_agent}" 51 | 52 | return agent 53 | 54 | 55 | class AlloyDBClient: 56 | def __init__( 57 | self, 58 | alloydb_api_endpoint: str, 59 | quota_project: Optional[str], 60 | credentials: Credentials, 61 | client: Optional[v1beta.AlloyDBAdminAsyncClient] = None, 62 | driver: Optional[str] = None, 63 | user_agent: Optional[str] = None, 64 | ) -> None: 65 | """ 66 | Establish the client to be used for AlloyDB API requests. 67 | 68 | Args: 69 | alloydb_api_endpoint (str): Base URL to use when calling 70 | the AlloyDB API endpoint. 71 | quota_project (str): The Project ID for an existing Google Cloud 72 | project. The project specified is used for quota and 73 | billing purposes. 74 | credentials (google.auth.credentials.Credentials): 75 | A credentials object created from the google-auth Python library. 76 | Must have the AlloyDB Admin scopes. For more info check out 77 | https://google-auth.readthedocs.io/en/latest/. 78 | client (v1beta.AlloyDBAdminAsyncClient): Async client used to make 79 | requests to AlloyDB APIs. 80 | Optional, defaults to None and creates new client. 81 | driver (str): Database driver to be used by the client. 82 | user_agent (str): The custom user-agent string to use in the HTTP 83 | header when making requests to AlloyDB APIs. 84 | Optional, defaults to None and uses a pre-defined one. 85 | """ 86 | user_agent = _format_user_agent(driver, user_agent) 87 | 88 | # TODO(rhatgadkar-goog): Rollback the PR of deciding between creating 89 | # AlloyDBAdminClient or AlloyDBAdminAsyncClient when either 90 | # https://github.com/grpc/grpc/issues/25364 is resolved or an async REST 91 | # transport for AlloyDBAdminAsyncClient gets introduced. 92 | # The issue is that the async gRPC transport does not work with multiple 93 | # event loops in the same process. So all calls to the AlloyDB Admin 94 | # API, even from multiple threads, need to be made to a single-event 95 | # loop. See https://github.com/GoogleCloudPlatform/alloydb-python-connector/issues/435 96 | # for more details. 97 | self._is_sync = False 98 | if client: 99 | self._client = client 100 | elif driver == "pg8000": 101 | self._client = v1beta.AlloyDBAdminClient( 102 | credentials=credentials, 103 | transport="grpc", 104 | client_options=ClientOptions( 105 | api_endpoint=alloydb_api_endpoint, 106 | quota_project_id=quota_project, 107 | ), 108 | client_info=ClientInfo( 109 | user_agent=user_agent, 110 | ), 111 | ) 112 | self._is_sync = True 113 | else: 114 | self._client = v1beta.AlloyDBAdminAsyncClient( 115 | credentials=credentials, 116 | transport="grpc_asyncio", 117 | client_options=ClientOptions( 118 | api_endpoint=alloydb_api_endpoint, 119 | quota_project_id=quota_project, 120 | ), 121 | client_info=ClientInfo( 122 | user_agent=user_agent, 123 | ), 124 | ) 125 | 126 | self._credentials = credentials 127 | # asyncpg does not currently support using metadata exchange 128 | # only use metadata exchange for pg8000 driver 129 | self._use_metadata = True if driver == "pg8000" else False 130 | self._user_agent = user_agent 131 | 132 | async def _get_metadata( 133 | self, 134 | project: str, 135 | region: str, 136 | cluster: str, 137 | name: str, 138 | ) -> dict[str, Optional[str]]: 139 | """ 140 | Fetch the metadata for a given AlloyDB instance. 141 | 142 | Call the AlloyDB APIs connectInfo method to retrieve the 143 | information about an AlloyDB instance that is used to create secure 144 | connections. 145 | 146 | Args: 147 | project (str): Google Cloud project ID that the AlloyDB instance 148 | resides in. 149 | region (str): Google Cloud region of the AlloyDB instance. 150 | cluster (str): The name of the AlloyDB cluster. 151 | name (str): The name of the AlloyDB instance. 152 | 153 | Returns: 154 | dict: IP addresses of the AlloyDB instance. 155 | """ 156 | parent = ( 157 | f"projects/{project}/locations/{region}/clusters/{cluster}/instances/{name}" 158 | ) 159 | 160 | req = v1beta.GetConnectionInfoRequest(parent=parent) 161 | if self._is_sync: 162 | resp = self._client.get_connection_info(request=req) 163 | else: 164 | resp = await self._client.get_connection_info(request=req) 165 | 166 | # Remove trailing period from PSC DNS name. 167 | psc_dns = resp.psc_dns_name 168 | if psc_dns: 169 | psc_dns = psc_dns.rstrip(".") 170 | 171 | return { 172 | "PRIVATE": resp.ip_address, 173 | "PUBLIC": resp.public_ip_address, 174 | "PSC": psc_dns, 175 | } 176 | 177 | async def _get_client_certificate( 178 | self, 179 | project: str, 180 | region: str, 181 | cluster: str, 182 | pub_key: str, 183 | ) -> tuple[str, list[str]]: 184 | """ 185 | Fetch a client certificate for the given AlloyDB cluster. 186 | 187 | Call the AlloyDB API's generateClientCertificate 188 | method to create a signed TLS certificate that is authorized to connect via the 189 | AlloyDB instance's serverside proxy. The cert is valid for twenty-four hours. 190 | 191 | Args: 192 | project (str): Google Cloud project ID that the AlloyDB instance 193 | resides in. 194 | region (str): Google Cloud region of the AlloyDB instance. 195 | cluster (str): The name of the AlloyDB cluster. 196 | pub_key (str): PEM-encoded client public key. 197 | 198 | Returns: 199 | tuple[str, list[str]]: tuple containing the CA certificate 200 | and certificate chain for the AlloyDB instance. 201 | """ 202 | parent = f"projects/{project}/locations/{region}/clusters/{cluster}" 203 | dur = duration_pb2.Duration() 204 | dur.seconds = 3600 205 | req = v1beta.GenerateClientCertificateRequest( 206 | parent=parent, 207 | cert_duration=dur, 208 | public_key=pub_key, 209 | use_metadata_exchange=self._use_metadata, 210 | ) 211 | if self._is_sync: 212 | resp = self._client.generate_client_certificate(request=req) 213 | else: 214 | resp = await self._client.generate_client_certificate(request=req) 215 | return (resp.ca_cert, resp.pem_certificate_chain) 216 | 217 | async def get_connection_info( 218 | self, 219 | project: str, 220 | region: str, 221 | cluster: str, 222 | name: str, 223 | keys: asyncio.Future, 224 | ) -> ConnectionInfo: 225 | """Immediately performs a full refresh operation using the AlloyDB API. 226 | 227 | Args: 228 | project (str): The name of the project the AlloyDB instance is 229 | located in. 230 | region (str): The region the AlloyDB instance is located in. 231 | cluster (str): The cluster the AlloyDB instance is located in. 232 | name (str): Name of the AlloyDB instance. 233 | keys (asyncio.Future): A future to the client's public-private key 234 | pair. 235 | 236 | Returns: 237 | ConnectionInfo: All the information required to connect securely to 238 | the AlloyDB instance. 239 | """ 240 | priv_key, pub_key = await keys 241 | 242 | # before making AlloyDB API calls, refresh creds if required 243 | if not self._credentials.token_state == TokenState.FRESH: 244 | self._credentials.refresh(requests.Request()) 245 | 246 | # fetch metadata 247 | metadata_task = asyncio.create_task( 248 | self._get_metadata( 249 | project, 250 | region, 251 | cluster, 252 | name, 253 | ) 254 | ) 255 | # generate client and CA certs 256 | certs_task = asyncio.create_task( 257 | self._get_client_certificate( 258 | project, 259 | region, 260 | cluster, 261 | pub_key, 262 | ) 263 | ) 264 | 265 | ip_addrs, certs = await asyncio.gather(metadata_task, certs_task) 266 | 267 | # unpack certs 268 | ca_cert, cert_chain = certs 269 | # get expiration from client certificate 270 | cert_obj = x509.load_pem_x509_certificate(cert_chain[0].encode("UTF-8")) 271 | expiration = cert_obj.not_valid_after_utc 272 | 273 | return ConnectionInfo( 274 | cert_chain, 275 | ca_cert, 276 | priv_key, 277 | ip_addrs, 278 | expiration, 279 | ) 280 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/connection_info.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from __future__ import annotations 16 | 17 | from dataclasses import dataclass 18 | import logging 19 | import ssl 20 | from typing import Optional, TYPE_CHECKING 21 | 22 | from aiofiles.tempfile import TemporaryDirectory 23 | 24 | from google.cloud.alloydbconnector.exceptions import IPTypeNotFoundError 25 | from google.cloud.alloydbconnector.utils import _write_to_file 26 | 27 | if TYPE_CHECKING: 28 | import datetime 29 | 30 | from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes 31 | 32 | from google.cloud.alloydbconnector.enums import IPTypes 33 | 34 | logger = logging.getLogger(name=__name__) 35 | 36 | 37 | @dataclass 38 | class ConnectionInfo: 39 | """Contains all necessary information to connect securely to the 40 | server-side Proxy running on an AlloyDB instance.""" 41 | 42 | cert_chain: list[str] 43 | ca_cert: str 44 | key: PrivateKeyTypes 45 | ip_addrs: dict[str, Optional[str]] 46 | expiration: datetime.datetime 47 | context: Optional[ssl.SSLContext] = None 48 | 49 | async def create_ssl_context(self) -> ssl.SSLContext: 50 | """Constructs a SSL/TLS context for the given connection info. 51 | 52 | Cache the SSL context to ensure we don't read from disk repeatedly when 53 | configuring a secure connection. 54 | """ 55 | # if SSL context is cached, use it 56 | if self.context is not None: 57 | return self.context 58 | 59 | # create TLS context 60 | context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) 61 | # force TLSv1.3 62 | context.minimum_version = ssl.TLSVersion.TLSv1_3 63 | 64 | # tmpdir and its contents are automatically deleted after the CA cert 65 | # and cert chain are loaded into the SSLcontext. The values 66 | # need to be written to files in order to be loaded by the SSLContext 67 | async with TemporaryDirectory() as tmpdir: 68 | ca_filename, cert_chain_filename, key_filename = await _write_to_file( 69 | tmpdir, self.ca_cert, self.cert_chain, self.key 70 | ) 71 | context.load_cert_chain(cert_chain_filename, keyfile=key_filename) 72 | context.load_verify_locations(cafile=ca_filename) 73 | # set class attribute to cache context for subsequent calls 74 | self.context = context 75 | return context 76 | 77 | def get_preferred_ip(self, ip_type: IPTypes) -> str: 78 | """Returns the first IP address for the instance, according to the preference 79 | supplied by ip_type. If no IP addressess with the given preference are found, 80 | an error is raised.""" 81 | ip_address = self.ip_addrs.get(ip_type.value) 82 | if ip_address is None: 83 | raise IPTypeNotFoundError( 84 | "AlloyDB instance does not have an IP addresses matching " 85 | f"type: '{ip_type.value}'" 86 | ) 87 | return ip_address 88 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/enums.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from __future__ import annotations 16 | 17 | from enum import Enum 18 | 19 | 20 | class IPTypes(Enum): 21 | """ 22 | Enum for specifying IP type to connect to AlloyDB with. 23 | """ 24 | 25 | PUBLIC = "PUBLIC" 26 | PRIVATE = "PRIVATE" 27 | PSC = "PSC" 28 | 29 | @classmethod 30 | def _missing_(cls, value: object) -> None: 31 | raise ValueError( 32 | f"Incorrect value for ip_type, got '{value}'. Want one of: " 33 | f"{', '.join([repr(m.value) for m in cls])}." 34 | ) 35 | 36 | 37 | class RefreshStrategy(Enum): 38 | """ 39 | Enum for specifying refresh strategy to connect to AlloyDB with. 40 | """ 41 | 42 | LAZY = "LAZY" 43 | BACKGROUND = "BACKGROUND" 44 | 45 | @classmethod 46 | def _missing_(cls, value: object) -> None: 47 | raise ValueError( 48 | f"Incorrect value for refresh_strategy, got '{value}'. Want one of: " 49 | f"{', '.join([repr(m.value) for m in cls])}." 50 | ) 51 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/exceptions.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | class RefreshError(Exception): 17 | pass 18 | 19 | 20 | class IPTypeNotFoundError(Exception): 21 | pass 22 | 23 | 24 | class ClosedConnectorError(Exception): 25 | pass 26 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/instance.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from __future__ import annotations 16 | 17 | import asyncio 18 | from datetime import datetime 19 | from datetime import timedelta 20 | from datetime import timezone 21 | import logging 22 | import re 23 | from typing import TYPE_CHECKING 24 | 25 | from google.cloud.alloydbconnector.connection_info import ConnectionInfo 26 | from google.cloud.alloydbconnector.exceptions import RefreshError 27 | from google.cloud.alloydbconnector.rate_limiter import AsyncRateLimiter 28 | from google.cloud.alloydbconnector.refresh_utils import _is_valid 29 | from google.cloud.alloydbconnector.refresh_utils import _seconds_until_refresh 30 | 31 | if TYPE_CHECKING: 32 | from cryptography.hazmat.primitives.asymmetric import rsa 33 | 34 | from google.cloud.alloydbconnector.client import AlloyDBClient 35 | 36 | logger = logging.getLogger(name=__name__) 37 | 38 | INSTANCE_URI_REGEX = re.compile( 39 | "projects/([^:]+(:[^:]+)?)/locations/([^:]+)/clusters/([^:]+)/instances/([^:]+)" 40 | ) 41 | 42 | 43 | def _parse_instance_uri(instance_uri: str) -> tuple[str, str, str, str]: 44 | # should take form "projects//locations//clusters//instances/" 45 | if INSTANCE_URI_REGEX.fullmatch(instance_uri) is None: 46 | raise ValueError( 47 | "Arg `instance_uri` must have " 48 | "format: projects//locations//clusters//instances/, projects/:/locations//clusters//instances/" 49 | f"got {instance_uri}." 50 | ) 51 | instance_uri_split = INSTANCE_URI_REGEX.split(instance_uri) 52 | return ( 53 | instance_uri_split[1], 54 | instance_uri_split[3], 55 | instance_uri_split[4], 56 | instance_uri_split[5], 57 | ) 58 | 59 | 60 | class RefreshAheadCache: 61 | """ 62 | Manages the information used to connect to the AlloyDB instance. 63 | 64 | Periodically calls the AlloyDB API, automatically refreshing the 65 | required information approximately 4 minutes before the previous 66 | certificate expires (every ~56 minutes). 67 | 68 | Args: 69 | instance_uri (str): The instance URI of the AlloyDB instance. 70 | ex. projects//locations//clusters//instances/ 71 | client (AlloyDBClient): Client used to make requests to AlloyDB APIs. 72 | keys (tuple[rsa.RSAPrivateKey, str]): Private and Public key pair. 73 | """ 74 | 75 | def __init__( 76 | self, 77 | instance_uri: str, 78 | client: AlloyDBClient, 79 | keys: asyncio.Future[tuple[rsa.RSAPrivateKey, str]], 80 | ) -> None: 81 | # validate and parse instance_uri 82 | self._project, self._region, self._cluster, self._name = _parse_instance_uri( 83 | instance_uri 84 | ) 85 | 86 | self._instance_uri = instance_uri 87 | self._client = client 88 | self._keys = keys 89 | self._refresh_rate_limiter = AsyncRateLimiter( 90 | max_capacity=2, 91 | rate=1 / 30, 92 | ) 93 | self._refresh_in_progress = asyncio.locks.Event() 94 | # For the initial refresh operation, set current = next so that 95 | # connection requests block until the first refresh is complete. 96 | self._current: asyncio.Task = self._schedule_refresh(0) 97 | self._next: asyncio.Task = self._current 98 | 99 | async def _perform_refresh(self) -> ConnectionInfo: 100 | """ 101 | Perform a refresh operation on an AlloyDB instance. 102 | 103 | Retrieves metadata and generates new client certificate 104 | required to connect securely to the AlloyDB instance. 105 | 106 | Returns: 107 | ConnectionInfo: Result of the refresh operation. 108 | """ 109 | self._refresh_in_progress.set() 110 | logger.debug( 111 | f"['{self._instance_uri}']: Connection info refresh operation started" 112 | ) 113 | 114 | try: 115 | await self._refresh_rate_limiter.acquire() 116 | connection_info = await self._client.get_connection_info( 117 | self._project, 118 | self._region, 119 | self._cluster, 120 | self._name, 121 | self._keys, 122 | ) 123 | logger.debug( 124 | f"['{self._instance_uri}']: Connection info refresh operation" 125 | " complete" 126 | ) 127 | logger.debug( 128 | f"['{self._instance_uri}']: Current certificate expiration = " 129 | f"{connection_info.expiration.isoformat()}" 130 | ) 131 | 132 | except Exception as e: 133 | logger.debug( 134 | f"['{self._instance_uri}']: Connection info refresh operation" 135 | f" failed: {str(e)}" 136 | ) 137 | raise 138 | 139 | finally: 140 | self._refresh_in_progress.clear() 141 | return connection_info 142 | 143 | def _schedule_refresh(self, delay: int) -> asyncio.Task: 144 | """ 145 | Schedule a refresh operation. 146 | 147 | Args: 148 | delay (int): Time in seconds to sleep before performing refresh. 149 | 150 | Returns: 151 | asyncio.Task[ConnectionInfo]: A task representing the scheduled 152 | refresh operation. 153 | """ 154 | return asyncio.create_task(self._refresh_operation(delay)) 155 | 156 | async def _refresh_operation(self, delay: int) -> ConnectionInfo: 157 | """ 158 | A coroutine that sleeps for the specified amount of time before 159 | running _perform_refresh. 160 | 161 | Args: 162 | delay (int): Time in seconds to sleep before performing refresh. 163 | 164 | Returns: 165 | ConnectionInfo: Refresh result for an AlloyDB instance. 166 | """ 167 | refresh_task: asyncio.Task 168 | try: 169 | if delay > 0: 170 | await asyncio.sleep(delay) 171 | refresh_task = asyncio.create_task(self._perform_refresh()) 172 | refresh_result = await refresh_task 173 | # check that refresh is valid 174 | if not await _is_valid(refresh_task): 175 | raise RefreshError( 176 | f"['{self._instance_uri}']: Invalid refresh operation. Certficate appears to be expired." 177 | ) 178 | except asyncio.CancelledError: 179 | logger.debug( 180 | f"['{self._instance_uri}']: Scheduled refresh operation cancelled" 181 | ) 182 | raise 183 | # bad refresh attempt 184 | except Exception: 185 | logger.info( 186 | f"['{self._instance_uri}']: " 187 | "An error occurred while performing refresh. " 188 | "Scheduling another refresh attempt immediately" 189 | ) 190 | # check if current refresh result is invalid (expired), 191 | # don't want to replace valid result with invalid refresh 192 | if not await _is_valid(self._current): 193 | self._current = refresh_task 194 | # schedule new refresh attempt immediately 195 | self._next = self._schedule_refresh(0) 196 | raise 197 | # if valid refresh, replace current with valid refresh result and schedule next refresh 198 | self._current = refresh_task 199 | # calculate refresh delay based on certificate expiration 200 | delay = _seconds_until_refresh(refresh_result.expiration) 201 | logger.debug( 202 | f"['{self._instance_uri}']: Connection info refresh operation" 203 | " scheduled for " 204 | f"{(datetime.now(timezone.utc) + timedelta(seconds=delay)).isoformat(timespec='seconds')} " 205 | f"(now + {timedelta(seconds=delay)})" 206 | ) 207 | self._next = self._schedule_refresh(delay) 208 | 209 | return refresh_result 210 | 211 | async def force_refresh(self) -> None: 212 | """ 213 | Schedules a new refresh operation immediately to be used 214 | for future connection attempts. 215 | """ 216 | # if next refresh is not already in progress, cancel it and schedule new one immediately 217 | if not self._refresh_in_progress.is_set(): 218 | self._next.cancel() 219 | self._next = self._schedule_refresh(0) 220 | # block all sequential connection attempts on the next refresh result if current is invalid 221 | if not await _is_valid(self._current): 222 | self._current = self._next 223 | 224 | async def connect_info(self) -> ConnectionInfo: 225 | """Retrieves ConnectionInfo instance for establishing a secure 226 | connection to the AlloyDB instance. 227 | """ 228 | return await self._current 229 | 230 | async def close(self) -> None: 231 | """ 232 | Cancel refresh tasks. 233 | """ 234 | logger.debug( 235 | f"['{self._instance_uri}']: Canceling connection info refresh" 236 | " operation tasks" 237 | ) 238 | self._current.cancel() 239 | self._next.cancel() 240 | # gracefully wait for tasks to cancel 241 | tasks = asyncio.gather(self._current, self._next, return_exceptions=True) 242 | await asyncio.wait_for(tasks, timeout=2.0) 243 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/lazy.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import asyncio 16 | from datetime import datetime 17 | from datetime import timedelta 18 | from datetime import timezone 19 | import logging 20 | from typing import Optional 21 | 22 | from google.cloud.alloydbconnector.client import AlloyDBClient 23 | from google.cloud.alloydbconnector.connection_info import ConnectionInfo 24 | from google.cloud.alloydbconnector.instance import _parse_instance_uri 25 | from google.cloud.alloydbconnector.refresh_utils import _refresh_buffer 26 | 27 | logger = logging.getLogger(name=__name__) 28 | 29 | 30 | class LazyRefreshCache: 31 | """Cache that refreshes connection info when a caller requests a connection. 32 | 33 | Only refreshes the cache when a new connection is requested and the current 34 | certificate is close to or already expired. 35 | 36 | This is the recommended option for serverless environments. 37 | """ 38 | 39 | def __init__( 40 | self, 41 | instance_uri: str, 42 | client: AlloyDBClient, 43 | keys: asyncio.Future, 44 | ) -> None: 45 | """Initializes a LazyRefreshCache instance. 46 | 47 | Args: 48 | instance_connection_string (str): The AlloyDB Instance's 49 | connection URI. 50 | client (AlloyDBClient): The AlloyDB client instance. 51 | keys (asyncio.Future): A future to the client's public-private key 52 | pair. 53 | """ 54 | # validate and parse instance connection name 55 | self._project, self._region, self._cluster, self._name = _parse_instance_uri( 56 | instance_uri 57 | ) 58 | self._instance_uri = instance_uri 59 | 60 | self._keys = keys 61 | self._client = client 62 | self._lock = asyncio.Lock() 63 | self._cached: Optional[ConnectionInfo] = None 64 | self._needs_refresh = False 65 | 66 | async def force_refresh(self) -> None: 67 | """ 68 | Invalidates the cache and configures the next call to 69 | connect_info() to retrieve a fresh ConnectionInfo instance. 70 | """ 71 | async with self._lock: 72 | self._needs_refresh = True 73 | 74 | async def connect_info(self) -> ConnectionInfo: 75 | """Retrieves ConnectionInfo instance for establishing a secure 76 | connection to the AlloyDB instance. 77 | """ 78 | async with self._lock: 79 | # If connection info is cached, check expiration. 80 | # Pad expiration with a buffer to give the client plenty of time to 81 | # establish a connection to the server with the certificate. 82 | if ( 83 | self._cached 84 | and not self._needs_refresh 85 | and datetime.now(timezone.utc) 86 | < (self._cached.expiration - timedelta(seconds=_refresh_buffer)) 87 | ): 88 | logger.debug( 89 | f"['{self._instance_uri}']: Connection info " 90 | "is still valid, using cached info" 91 | ) 92 | return self._cached 93 | logger.debug( 94 | f"['{self._instance_uri}']: Connection info " 95 | "refresh operation started" 96 | ) 97 | try: 98 | conn_info = await self._client.get_connection_info( 99 | self._project, 100 | self._region, 101 | self._cluster, 102 | self._name, 103 | self._keys, 104 | ) 105 | except Exception as e: 106 | logger.debug( 107 | f"['{self._instance_uri}']: Connection info " 108 | f"refresh operation failed: {str(e)}" 109 | ) 110 | raise 111 | logger.debug( 112 | f"['{self._instance_uri}']: Connection info " 113 | "refresh operation completed successfully" 114 | ) 115 | logger.debug( 116 | f"['{self._instance_uri}']: Current certificate " 117 | f"expiration = {str(conn_info.expiration)}" 118 | ) 119 | self._cached = conn_info 120 | self._needs_refresh = False 121 | return conn_info 122 | 123 | async def close(self) -> None: 124 | """Close is a no-op and provided purely for a consistent interface with 125 | other cache types. 126 | """ 127 | pass 128 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/pg8000.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from typing import Any, TYPE_CHECKING 15 | 16 | if TYPE_CHECKING: 17 | import ssl 18 | 19 | import pg8000 20 | 21 | 22 | def connect(sock: "ssl.SSLSocket", **kwargs: Any) -> "pg8000.dbapi.Connection": 23 | """Create a pg8000 DBAPI connection object. 24 | 25 | Args: 26 | sock (ssl.SSLSocket): SSL/TLS secure socket stream connected to the 27 | AlloyDB proxy server. 28 | 29 | Returns: 30 | pg8000.dbapi.Connection: A pg8000 Connection object for 31 | the AlloyDB instance. 32 | """ 33 | try: 34 | import pg8000 35 | except ImportError: 36 | raise ImportError( 37 | 'Unable to import module "pg8000." Please install and try again.' 38 | ) 39 | 40 | user = kwargs.pop("user") 41 | db = kwargs.pop("db") 42 | passwd = kwargs.pop("password", None) 43 | return pg8000.dbapi.connect( 44 | user, 45 | database=db, 46 | password=passwd, 47 | sock=sock, 48 | **kwargs, 49 | ) 50 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GoogleCloudPlatform/alloydb-python-connector/86d38bbff1c409cae84696a81bb009a71a3e2a9d/google/cloud/alloydbconnector/py.typed -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/rate_limiter.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import asyncio 16 | 17 | 18 | class AsyncRateLimiter(object): 19 | """ 20 | An asyncio-compatible rate limiter. 21 | 22 | Uses the Token Bucket algorithm (https://en.wikipedia.org/wiki/Token_bucket) 23 | to limit the number of function calls over a time interval using an event queue. 24 | 25 | Args: 26 | max_capacity (int): The maximum capacity of tokens the bucket 27 | will store at any one time. Defaults to 1. 28 | 29 | rate (float): The number of tokens that should be added per second. 30 | Defaults to 1 / 60. 31 | """ 32 | 33 | def __init__( 34 | self, 35 | max_capacity: int = 1, 36 | rate: float = 1 / 60, 37 | ) -> None: 38 | self._rate = rate 39 | self._max_capacity = max_capacity 40 | self._loop = asyncio.get_running_loop() 41 | self._tokens: float = max_capacity 42 | self._last_token_update = self._loop.time() 43 | self._lock = asyncio.Lock() 44 | 45 | def _update_token_count(self) -> None: 46 | """ 47 | Calculates how much time has passed since the last leak and removes the 48 | appropriate amount of events from the queue. 49 | Leaking is done lazily, meaning that if there is a large time gap between 50 | leaks, the next set of calls might be a burst if burst_size > 1 51 | """ 52 | now = self._loop.time() 53 | time_elapsed = now - self._last_token_update 54 | new_tokens = time_elapsed * self._rate 55 | self._tokens = min(new_tokens + self._tokens, self._max_capacity) 56 | self._last_token_update = now 57 | 58 | async def _wait_for_next_token(self) -> None: 59 | """ 60 | Wait until enough time has elapsed to add another token. 61 | """ 62 | token_deficit = 1 - self._tokens 63 | if token_deficit > 0: 64 | wait_time = token_deficit / self._rate 65 | await asyncio.sleep(wait_time) 66 | 67 | async def acquire(self) -> None: 68 | """ 69 | Waits for a token to become available, if necessary, then subtracts token and allows 70 | request to go through. 71 | """ 72 | async with self._lock: 73 | self._update_token_count() 74 | if self._tokens < 1: 75 | await self._wait_for_next_token() 76 | self._update_token_count() 77 | self._tokens -= 1 78 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/refresh_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from __future__ import annotations 16 | 17 | import asyncio 18 | from datetime import datetime 19 | from datetime import timezone 20 | import logging 21 | 22 | logger = logging.getLogger(name=__name__) 23 | 24 | # _refresh_buffer is the amount of time before a refresh's result expires 25 | # that a new refresh operation begins. 26 | _refresh_buffer: int = 4 * 60 # 4 minutes 27 | 28 | 29 | def _seconds_until_refresh(expiration: datetime) -> int: 30 | """ 31 | Calculates the duration to wait before starting the next refresh. 32 | Usually the duration will be half of the time until certificate 33 | expiration. 34 | 35 | Args: 36 | expiration (datetime.datetime): Time of certificate expiration. 37 | Returns: 38 | int: Time in seconds to wait before performing next refresh. 39 | """ 40 | 41 | duration = int((expiration - datetime.now(timezone.utc)).total_seconds()) 42 | 43 | # if certificate duration is less than 1 hour 44 | if duration < 3600: 45 | # something is wrong with certificate, refresh now 46 | if duration < _refresh_buffer: 47 | return 0 48 | # otherwise wait until 4 minutes before expiration for next refresh 49 | return duration - _refresh_buffer 50 | return duration // 2 51 | 52 | 53 | async def _is_valid(task: asyncio.Task) -> bool: 54 | try: 55 | result = await task 56 | # valid if current time is before cert expiration 57 | if datetime.now(timezone.utc) < result.expiration: 58 | return True 59 | except Exception: 60 | # suppress any errors from task 61 | logger.debug("Current refresh result is invalid.") 62 | return False 63 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/static.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from datetime import datetime 16 | from datetime import timedelta 17 | from datetime import timezone 18 | import io 19 | import json 20 | 21 | from cryptography.hazmat.primitives import serialization 22 | 23 | from google.cloud.alloydbconnector.connection_info import ConnectionInfo 24 | 25 | 26 | class StaticConnectionInfoCache: 27 | """ 28 | StaticConnectionInfoCache creates a connection info cache that will always 29 | return a pre-defined connection info. This is a *dev-only* option and 30 | should not be used in production as it will result in failed connections 31 | after the client certificate expires. It is also subject to breaking changes 32 | in the format. NOTE: The static connection info is not refreshed by the 33 | connector. The JSON format supports multiple instances, regardless of 34 | cluster. 35 | 36 | This static connection info should hold JSON with the following format: 37 | { 38 | "publicKey": "", 39 | "privateKey": "", 40 | "projects//locations//clusters//instances/": { 41 | "ipAddress": "", 42 | "publicIpAddress": "", 43 | "pscInstanceConfig": { 44 | "pscDnsName": "" 45 | }, 46 | "pemCertificateChain": [ 47 | "", "", "" 48 | ], 49 | "caCert": "" 50 | } 51 | } 52 | """ 53 | 54 | def __init__(self, instance_uri: str, static_conn_info: io.TextIOBase) -> None: 55 | """ 56 | Initializes a StaticConnectionInfoCache instance. 57 | 58 | Args: 59 | instance_uri (str): The AlloyDB instance's connection URI. 60 | static_conn_info (io.TextIOBase): The static connection info JSON. 61 | """ 62 | static_info = json.load(static_conn_info) 63 | ca_cert = static_info[instance_uri]["caCert"] 64 | cert_chain = static_info[instance_uri]["pemCertificateChain"] 65 | dns = "" 66 | if static_info[instance_uri]["pscInstanceConfig"]: 67 | dns = static_info[instance_uri]["pscInstanceConfig"]["pscDnsName"].rstrip( 68 | "." 69 | ) 70 | ip_addrs = { 71 | "PRIVATE": static_info[instance_uri]["ipAddress"], 72 | "PUBLIC": static_info[instance_uri]["publicIpAddress"], 73 | "PSC": dns, 74 | } 75 | expiration = datetime.now(timezone.utc) + timedelta(hours=1) 76 | priv_key = static_info["privateKey"] 77 | priv_key_bytes = serialization.load_pem_private_key( 78 | priv_key.encode("UTF-8"), password=None 79 | ) 80 | self._info = ConnectionInfo( 81 | cert_chain, ca_cert, priv_key_bytes, ip_addrs, expiration 82 | ) 83 | 84 | async def force_refresh(self) -> None: 85 | """ 86 | This is a no-op as the cache holds only static connection information 87 | and does no refresh. 88 | """ 89 | pass 90 | 91 | async def connect_info(self) -> ConnectionInfo: 92 | """ 93 | Retrieves ConnectionInfo instance for establishing a secure 94 | connection to the AlloyDB instance. 95 | """ 96 | return self._info 97 | 98 | async def close(self) -> None: 99 | """ 100 | This is a no-op. 101 | """ 102 | pass 103 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import typing 16 | 17 | from google.cloud.alloydbconnector.instance import RefreshAheadCache 18 | from google.cloud.alloydbconnector.lazy import LazyRefreshCache 19 | from google.cloud.alloydbconnector.static import StaticConnectionInfoCache 20 | 21 | CacheTypes = typing.Union[ 22 | RefreshAheadCache, LazyRefreshCache, StaticConnectionInfoCache 23 | ] 24 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from __future__ import annotations 16 | 17 | import re 18 | 19 | import aiofiles 20 | from cryptography.hazmat.primitives import serialization 21 | from cryptography.hazmat.primitives.asymmetric import rsa 22 | from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes 23 | 24 | 25 | async def _write_to_file( 26 | dir_path: str, ca_cert: str, cert_chain: list[str], key: PrivateKeyTypes 27 | ) -> tuple[str, str, str]: 28 | """ 29 | Helper function to write the server_ca, client certificate and 30 | private key to .pem files in a given directory. 31 | """ 32 | ca_filename = f"{dir_path}/ca.pem" 33 | cert_chain_filename = f"{dir_path}/chain.pem" 34 | key_filename = f"{dir_path}/priv.pem" 35 | 36 | key_bytes = key.private_bytes( 37 | encoding=serialization.Encoding.PEM, 38 | format=serialization.PrivateFormat.TraditionalOpenSSL, 39 | encryption_algorithm=serialization.NoEncryption(), 40 | ) 41 | 42 | async with aiofiles.open(ca_filename, "w+") as ca_out: 43 | await ca_out.write(ca_cert) 44 | async with aiofiles.open(cert_chain_filename, "w+") as chain_out: 45 | await chain_out.write("".join(cert_chain)) 46 | async with aiofiles.open(key_filename, "wb") as priv_out: 47 | await priv_out.write(key_bytes) 48 | 49 | return (ca_filename, cert_chain_filename, key_filename) 50 | 51 | 52 | async def generate_keys() -> tuple[rsa.RSAPrivateKey, str]: 53 | priv_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) 54 | pub_key = ( 55 | priv_key.public_key() 56 | .public_bytes( 57 | encoding=serialization.Encoding.PEM, 58 | format=serialization.PublicFormat.SubjectPublicKeyInfo, 59 | ) 60 | .decode("UTF-8") 61 | ) 62 | return (priv_key, pub_key) 63 | 64 | 65 | def strip_http_prefix(url: str) -> str: 66 | """ 67 | Returns a new URL with 'http://' or 'https://' prefix removed. 68 | """ 69 | m = re.search(r"^(https?://)?(.+)", url) 70 | if m is None: 71 | return "" 72 | return m.group(2) 73 | -------------------------------------------------------------------------------- /google/cloud/alloydbconnector/version.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | __version__ = "1.9.0" 16 | -------------------------------------------------------------------------------- /noxfile.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from __future__ import absolute_import 16 | 17 | import os 18 | 19 | import nox 20 | 21 | BLACK_VERSION = "black==23.12.1" 22 | ISORT_VERSION = "isort==5.13.2" 23 | LINT_PATHS = ["google", "tests", "noxfile.py"] 24 | 25 | SYSTEM_TEST_PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.12", "3.13"] 26 | UNIT_TEST_PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.12", "3.13"] 27 | 28 | 29 | @nox.session 30 | def lint(session): 31 | """Run linters. 32 | 33 | Returns a failure if the linters find linting errors or sufficiently 34 | serious code quality issues. 35 | """ 36 | session.install("-r", "requirements-test.txt") 37 | session.install("-r", "requirements.txt") 38 | session.install( 39 | "flake8", 40 | "flake8-annotations", 41 | "mypy", 42 | BLACK_VERSION, 43 | ISORT_VERSION, 44 | "build", 45 | "twine", 46 | ) 47 | session.run( 48 | "isort", 49 | "--fss", 50 | "--check-only", 51 | "--diff", 52 | "--profile=google", 53 | *LINT_PATHS, 54 | ) 55 | session.run( 56 | "black", 57 | "--check", 58 | "--diff", 59 | *LINT_PATHS, 60 | ) 61 | session.run( 62 | "flake8", 63 | "google", 64 | "tests", 65 | ) 66 | session.run( 67 | "mypy", 68 | "-p", 69 | "google", 70 | "--install-types", 71 | "--non-interactive", 72 | "--show-traceback", 73 | ) 74 | # verify that pyproject.toml is valid 75 | session.run("python", "-m", "build", "--sdist") 76 | session.run("twine", "check", "--strict", "dist/*") 77 | 78 | 79 | @nox.session 80 | def blacken(session): 81 | """Run black. 82 | 83 | Format code to uniform standard. 84 | """ 85 | session.install(BLACK_VERSION) 86 | session.run( 87 | "black", 88 | *LINT_PATHS, 89 | ) 90 | 91 | 92 | @nox.session() 93 | def format(session): 94 | """ 95 | Run isort to sort imports. Then run black 96 | to format code to uniform standard. 97 | """ 98 | session.install(BLACK_VERSION, ISORT_VERSION) 99 | # Use the --fss option to sort imports using strict alphabetical order. 100 | # See https://pycqa.github.io/isort/docs/configuration/options.html#force-sort-within-sectionss 101 | session.run( 102 | "isort", 103 | "--fss", 104 | "--profile=google", 105 | *LINT_PATHS, 106 | ) 107 | session.run( 108 | "black", 109 | *LINT_PATHS, 110 | ) 111 | 112 | 113 | @nox.session() 114 | def cover(session): 115 | """Run the final coverage report. 116 | 117 | This outputs the coverage report aggregating coverage from the unit 118 | test runs (not system test runs), and then erases coverage data. 119 | """ 120 | session.install("coverage", "pytest-cov") 121 | session.run("coverage", "report", "--show-missing", "--fail-under=0") 122 | 123 | session.run("coverage", "erase") 124 | 125 | 126 | def default(session, path): 127 | # Install all test dependencies, then install this package in-place. 128 | session.install("-r", "requirements-test.txt") 129 | session.install(".") 130 | session.install("-r", "requirements.txt") 131 | # Run pytest with coverage. 132 | # Using the coverage command instead of `pytest --cov`, because 133 | # `pytest ---cov` causes the module to be initialized twice, which returns 134 | # this error: "ImportError: PyO3 modules compiled for CPython 3.8 or older 135 | # may only be initialized once per interpreter process". More info about 136 | # this is stated here: https://github.com/pytest-dev/pytest-cov/issues/614. 137 | session.run( 138 | "coverage", 139 | "run", 140 | "--include=*/google/cloud/alloydbconnector/*.py", 141 | "-m", 142 | "pytest", 143 | "-v", 144 | path, 145 | *session.posargs, 146 | ) 147 | session.run( 148 | "coverage", 149 | "xml", 150 | "-o", 151 | "sponge_log.xml", 152 | ) 153 | 154 | 155 | @nox.session(python=UNIT_TEST_PYTHON_VERSIONS) 156 | def unit(session): 157 | """Run the unit test suite.""" 158 | default(session, os.path.join("tests", "unit")) 159 | 160 | 161 | @nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) 162 | def system(session): 163 | default(session, os.path.join("tests", "system")) 164 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | [build-system] 16 | requires = ["setuptools"] 17 | build-backend = "setuptools.build_meta" 18 | 19 | [project] 20 | description = "A Python client library for connecting securely to your Google Cloud AlloyDB instances." 21 | name = "google-cloud-alloydb-connector" 22 | authors = [{ name = "Google LLC", email = "googleapis-packages@google.com" }] 23 | license = "Apache-2.0" 24 | license-files = ["LICENSE"] 25 | requires-python = ">=3.9" 26 | readme = "README.md" 27 | classifiers = [ 28 | # Should be one of: 29 | # "Development Status :: 3 - Alpha" 30 | # "Development Status :: 4 - Beta" 31 | # "Development Status :: 5 - Production/Stable" 32 | "Development Status :: 5 - Production/Stable", 33 | "Intended Audience :: Developers", 34 | "Programming Language :: Python", 35 | "Programming Language :: Python :: 3", 36 | "Programming Language :: Python :: 3.9", 37 | "Programming Language :: Python :: 3.10", 38 | "Programming Language :: Python :: 3.11", 39 | "Programming Language :: Python :: 3.12", 40 | "Programming Language :: Python :: 3.13", 41 | "Operating System :: OS Independent", 42 | ] 43 | dependencies = [ 44 | "aiofiles", 45 | "aiohttp", 46 | "cryptography>=42.0.0", 47 | "requests", 48 | "google-auth", 49 | "protobuf", 50 | "google-cloud-alloydb", 51 | "google-api-core", 52 | ] 53 | dynamic = ["version"] 54 | 55 | [project.urls] 56 | Homepage = "https://github.com/GoogleCloudPlatform/alloydb-python-connector" 57 | Repository = "https://github.com/GoogleCloudPlatform/alloydb-python-connector.git" 58 | Issues = "https://github.com/GoogleCloudPlatform/alloydb-python-connector/issues" 59 | Changelog = "https://github.com/GoogleCloudPlatform/alloydb-python-connector/blob/main/CHANGELOG.md" 60 | 61 | [project.optional-dependencies] 62 | pg8000 = ["pg8000>=1.31.1"] 63 | asyncpg = ["asyncpg>=0.30.0"] 64 | 65 | [tool.setuptools.dynamic] 66 | version = { attr = "google.cloud.alloydbconnector.version.__version__" } 67 | 68 | [tool.setuptools.package-data] 69 | "google.cloud.alloydbconnector" = ["py.typed"] 70 | 71 | [tool.setuptools.packages.find] 72 | # Only include packages under the 'google' namespace. Do not include tests, 73 | # benchmarks, etc. 74 | include = ["google*"] 75 | 76 | [tool.mypy] 77 | python_version = "3.9" 78 | namespace_packages = true 79 | ignore_missing_imports = true 80 | warn_unused_configs = true 81 | exclude = ['docs/*'] 82 | 83 | [tool.pytest.ini_options] 84 | asyncio_mode = "auto" 85 | -------------------------------------------------------------------------------- /requirements-test.txt: -------------------------------------------------------------------------------- 1 | asyncpg==0.30.0 2 | mock==5.2.0 3 | pg8000==1.31.2 4 | 5 | # psycopg 2.9.10 isn't supported on Python 3.9 for macos-latest GitHub runner. 6 | # It is supported for macos-12 runner: 7 | # https://github.com/psycopg/psycopg2/issues/1737. But macos-12 runner is 8 | # deprecated: https://github.com/actions/runner-images/issues/10721. So we 9 | # install psycopg 2.9.9 on Python 3.9 for macos-latest runner. 10 | psycopg2-binary==2.9.9; python_version == "3.9" and sys_platform == "darwin" 11 | psycopg2-binary==2.9.10; python_version != "3.9" or sys_platform != "darwin" 12 | 13 | pytest==8.4.1 14 | pytest-asyncio==1.0.0 15 | pytest-cov==6.2.1 16 | pytest-aiohttp==1.1.0 17 | SQLAlchemy[asyncio]==2.0.41 18 | aioresponses==0.7.8 19 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | aiofiles==24.1.0 2 | aiohttp==3.12.13 3 | cryptography==45.0.5 4 | google-auth==2.40.3 5 | requests==2.32.4 6 | protobuf==6.31.1 7 | -------------------------------------------------------------------------------- /tests/system/test_alloydb_connector_package.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import google.cloud.alloydb.connector as conn_old 16 | import google.cloud.alloydbconnector as conn_new 17 | 18 | 19 | def test_alloydbconnector_package() -> None: 20 | """ 21 | Test imported objects are same in google.cloud.alloydb.connector and 22 | google.cloud.alloydbconnector packages. 23 | """ 24 | assert conn_old.AsyncConnector == conn_new.AsyncConnector 25 | assert conn_old.Connector == conn_new.Connector 26 | assert conn_old.IPTypes == conn_new.IPTypes 27 | assert conn_old.RefreshStrategy == conn_new.RefreshStrategy 28 | assert conn_old.__version__ == conn_new.__version__ 29 | -------------------------------------------------------------------------------- /tests/system/test_asyncpg_connection.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | 17 | # [START alloydb_sqlalchemy_connect_async_connector] 18 | import asyncpg 19 | import sqlalchemy 20 | import sqlalchemy.ext.asyncio 21 | 22 | from google.cloud.alloydbconnector import AsyncConnector 23 | 24 | 25 | async def create_sqlalchemy_engine( 26 | inst_uri: str, 27 | user: str, 28 | password: str, 29 | db: str, 30 | refresh_strategy: str = "background", 31 | ) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, AsyncConnector]: 32 | """Creates a connection pool for an AlloyDB instance and returns the pool 33 | and the connector. Callers are responsible for closing the pool and the 34 | connector. 35 | 36 | A sample invocation looks like: 37 | 38 | engine, connector = await create_sqlalchemy_engine( 39 | inst_uri, 40 | user, 41 | password, 42 | db, 43 | ) 44 | async with engine.connect() as conn: 45 | time = await conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 46 | curr_time = time[0] 47 | # do something with query result 48 | await connector.close() 49 | 50 | Args: 51 | instance_uri (str): 52 | The instance URI specifies the instance relative to the project, 53 | region, and cluster. For example: 54 | "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" 55 | user (str): 56 | The database user name, e.g., postgres 57 | password (str): 58 | The database user's password, e.g., secret-password 59 | db (str): 60 | The name of the database, e.g., mydb 61 | refresh_strategy (Optional[str]): 62 | Refresh strategy for the AlloyDB Connector. Can be one of "lazy" 63 | or "background". For serverless environments use "lazy" to avoid 64 | errors resulting from CPU being throttled. 65 | """ 66 | connector = AsyncConnector(refresh_strategy=refresh_strategy) 67 | 68 | # create SQLAlchemy connection pool 69 | engine = sqlalchemy.ext.asyncio.create_async_engine( 70 | "postgresql+asyncpg://", 71 | async_creator=lambda: connector.connect( 72 | inst_uri, 73 | "asyncpg", 74 | user=user, 75 | password=password, 76 | db=db, 77 | ), 78 | execution_options={"isolation_level": "AUTOCOMMIT"}, 79 | ) 80 | return engine, connector 81 | 82 | 83 | # [END alloydb_sqlalchemy_connect_async_connector] 84 | 85 | 86 | async def create_asyncpg_pool( 87 | instance_connection_name: str, 88 | user: str, 89 | password: str, 90 | db: str, 91 | refresh_strategy: str = "background", 92 | ) -> tuple[asyncpg.Pool, AsyncConnector]: 93 | """Creates a native asyncpg connection pool for an AlloyDB instance and 94 | returns the pool and the connector. Callers are responsible for closing the 95 | pool and the connector. 96 | 97 | A sample invocation looks like: 98 | 99 | pool, connector = await create_asyncpg_pool( 100 | inst_conn_name, 101 | user, 102 | password, 103 | db, 104 | ) 105 | async with pool.acquire() as conn: 106 | hello = await conn.fetch("SELECT 'Hello World!'") 107 | # do something with query result 108 | await connector.close() 109 | 110 | Args: 111 | instance_connection_name (str): 112 | The instance connection name specifies the instance relative to the 113 | project and region. For example: "my-project:my-region:my-instance" 114 | user (str): 115 | The database user name, e.g., postgres 116 | password (str): 117 | The database user's password, e.g., secret-password 118 | db (str): 119 | The name of the database, e.g., mydb 120 | refresh_strategy (Optional[str]): 121 | Refresh strategy for the Cloud SQL Connector. Can be one of "lazy" 122 | or "background". For serverless environments use "lazy" to avoid 123 | errors resulting from CPU being throttled. 124 | """ 125 | connector = AsyncConnector(refresh_strategy=refresh_strategy) 126 | 127 | # create native asyncpg pool (requires asyncpg version >=0.30.0) 128 | pool = await asyncpg.create_pool( 129 | instance_connection_name, 130 | connect=lambda instance_connection_name, **kwargs: connector.connect( 131 | instance_connection_name, 132 | "asyncpg", 133 | user=user, 134 | password=password, 135 | db=db, 136 | ), 137 | ) 138 | return pool, connector 139 | 140 | 141 | async def test_sqlalchemy_connection_with_asyncpg() -> None: 142 | """Basic test to get time from database.""" 143 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 144 | user = os.environ["ALLOYDB_USER"] 145 | password = os.environ["ALLOYDB_PASS"] 146 | db = os.environ["ALLOYDB_DB"] 147 | 148 | pool, connector = await create_sqlalchemy_engine(inst_uri, user, password, db) 149 | 150 | async with pool.connect() as conn: 151 | res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone() 152 | assert res[0] == 1 153 | 154 | await connector.close() 155 | 156 | 157 | async def test_lazy_sqlalchemy_connection_with_asyncpg() -> None: 158 | """Basic test to get time from database.""" 159 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 160 | user = os.environ["ALLOYDB_USER"] 161 | password = os.environ["ALLOYDB_PASS"] 162 | db = os.environ["ALLOYDB_DB"] 163 | 164 | pool, connector = await create_sqlalchemy_engine( 165 | inst_uri, user, password, db, "lazy" 166 | ) 167 | 168 | async with pool.connect() as conn: 169 | res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone() 170 | assert res[0] == 1 171 | 172 | await connector.close() 173 | 174 | 175 | async def test_connection_with_asyncpg() -> None: 176 | """Basic test to get time from database.""" 177 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 178 | user = os.environ["ALLOYDB_USER"] 179 | password = os.environ["ALLOYDB_PASS"] 180 | db = os.environ["ALLOYDB_DB"] 181 | 182 | pool, connector = await create_asyncpg_pool(inst_uri, user, password, db) 183 | 184 | async with pool.acquire() as conn: 185 | res = await conn.fetch("SELECT 1") 186 | assert res[0][0] == 1 187 | 188 | await connector.close() 189 | 190 | 191 | async def test_lazy_connection_with_asyncpg() -> None: 192 | """Basic test to get time from database.""" 193 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 194 | user = os.environ["ALLOYDB_USER"] 195 | password = os.environ["ALLOYDB_PASS"] 196 | db = os.environ["ALLOYDB_DB"] 197 | 198 | pool, connector = await create_asyncpg_pool(inst_uri, user, password, db, "lazy") 199 | 200 | async with pool.acquire() as conn: 201 | res = await conn.fetch("SELECT 1") 202 | assert res[0][0] == 1 203 | 204 | await connector.close() 205 | -------------------------------------------------------------------------------- /tests/system/test_asyncpg_iam_authn.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from datetime import datetime 16 | import os 17 | 18 | # [START alloydb_sqlalchemy_connect_async_connector_iam_authn] 19 | import sqlalchemy 20 | import sqlalchemy.ext.asyncio 21 | 22 | from google.cloud.alloydbconnector import AsyncConnector 23 | 24 | 25 | async def create_sqlalchemy_engine( 26 | inst_uri: str, user: str, db: str, refresh_strategy: str = "background" 27 | ) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, AsyncConnector]: 28 | """Creates a connection pool for an AlloyDB instance and returns the pool 29 | and the connector. Callers are responsible for closing the pool and the 30 | connector. 31 | 32 | A sample invocation looks like: 33 | 34 | pool, connector = await create_sqlalchemy_engine( 35 | inst_uri, 36 | user, 37 | db, 38 | ) 39 | async with pool.connect() as conn: 40 | time = (await conn.execute(sqlalchemy.text("SELECT NOW()"))).fetchone() 41 | conn.commit() 42 | curr_time = time[0] 43 | # do something with query result 44 | await connector.close() 45 | 46 | Args: 47 | instance_uri (str): 48 | The instance URI specifies the instance relative to the project, 49 | region, and cluster. For example: 50 | "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" 51 | user (str): 52 | The formatted IAM database username. 53 | e.g., my-email@test.com, service-account@project-id.iam 54 | db (str): 55 | The name of the database, e.g., mydb 56 | refresh_strategy (Optional[str]): 57 | Refresh strategy for the AlloyDB Connector. Can be one of "lazy" 58 | or "background". For serverless environments use "lazy" to avoid 59 | errors resulting from CPU being throttled. 60 | """ 61 | connector = AsyncConnector(refresh_strategy=refresh_strategy) 62 | 63 | # create async SQLAlchemy connection pool 64 | engine = sqlalchemy.ext.asyncio.create_async_engine( 65 | "postgresql+asyncpg://", 66 | async_creator=lambda: connector.connect( 67 | inst_uri, 68 | "asyncpg", 69 | user=user, 70 | db=db, 71 | enable_iam_auth=True, 72 | ), 73 | execution_options={"isolation_level": "AUTOCOMMIT"}, 74 | ) 75 | return engine, connector 76 | 77 | 78 | # [END alloydb_sqlalchemy_connect_async_connector_iam_authn] 79 | 80 | 81 | async def test_asyncpg_iam_authn_time() -> None: 82 | """Basic test to get time from database.""" 83 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 84 | user = os.environ["ALLOYDB_IAM_USER"] 85 | db = os.environ["ALLOYDB_DB"] 86 | 87 | pool, connector = await create_sqlalchemy_engine(inst_uri, user, db) 88 | async with pool.connect() as conn: 89 | time = (await conn.execute(sqlalchemy.text("SELECT NOW()"))).fetchone() 90 | curr_time = time[0] 91 | assert type(curr_time) is datetime 92 | await connector.close() 93 | # cleanup AsyncEngine 94 | await pool.dispose() 95 | 96 | 97 | async def test_asyncpg_iam_authn_lazy() -> None: 98 | """Basic test to get time from database.""" 99 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 100 | user = os.environ["ALLOYDB_IAM_USER"] 101 | db = os.environ["ALLOYDB_DB"] 102 | 103 | pool, connector = await create_sqlalchemy_engine(inst_uri, user, db, "lazy") 104 | async with pool.connect() as conn: 105 | time = (await conn.execute(sqlalchemy.text("SELECT NOW()"))).fetchone() 106 | curr_time = time[0] 107 | assert type(curr_time) is datetime 108 | await connector.close() 109 | # cleanup AsyncEngine 110 | await pool.dispose() 111 | -------------------------------------------------------------------------------- /tests/system/test_asyncpg_psc.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | 17 | import pytest 18 | import sqlalchemy 19 | import sqlalchemy.ext.asyncio 20 | 21 | from google.cloud.alloydbconnector import AsyncConnector 22 | 23 | 24 | async def create_sqlalchemy_engine( 25 | inst_uri: str, 26 | user: str, 27 | password: str, 28 | db: str, 29 | ) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, AsyncConnector]: 30 | """Creates a connection pool for an AlloyDB instance and returns the pool 31 | and the connector. Callers are responsible for closing the pool and the 32 | connector. 33 | 34 | A sample invocation looks like: 35 | 36 | engine, connector = await create_sqlalchemy_engine( 37 | inst_uri, 38 | user, 39 | password, 40 | db, 41 | ) 42 | async with engine.connect() as conn: 43 | time = await conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 44 | curr_time = time[0] 45 | # do something with query result 46 | await connector.close() 47 | 48 | Args: 49 | instance_uri (str): 50 | The instance URI specifies the instance relative to the project, 51 | region, and cluster. For example: 52 | "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" 53 | user (str): 54 | The database user name, e.g., postgres 55 | password (str): 56 | The database user's password, e.g., secret-password 57 | db_name (str): 58 | The name of the database, e.g., mydb 59 | """ 60 | connector = AsyncConnector() 61 | 62 | # create SQLAlchemy connection pool 63 | engine = sqlalchemy.ext.asyncio.create_async_engine( 64 | "postgresql+asyncpg://", 65 | async_creator=lambda: connector.connect( 66 | inst_uri, 67 | "asyncpg", 68 | user=user, 69 | password=password, 70 | db=db, 71 | ip_type="PSC", 72 | ), 73 | execution_options={"isolation_level": "AUTOCOMMIT"}, 74 | ) 75 | return engine, connector 76 | 77 | 78 | @pytest.mark.asyncio 79 | async def test_connection_with_asyncpg() -> None: 80 | """Basic test to get time from database.""" 81 | inst_uri = os.environ["ALLOYDB_PSC_INSTANCE_URI"] 82 | user = os.environ["ALLOYDB_USER"] 83 | password = os.environ["ALLOYDB_PASS"] 84 | db = os.environ["ALLOYDB_DB"] 85 | 86 | pool, connector = await create_sqlalchemy_engine(inst_uri, user, password, db) 87 | 88 | async with pool.connect() as conn: 89 | res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone() 90 | assert res[0] == 1 91 | 92 | await connector.close() 93 | -------------------------------------------------------------------------------- /tests/system/test_asyncpg_public_ip.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | 17 | # [START alloydb_sqlalchemy_connect_async_connector_public_ip] 18 | import pytest 19 | import sqlalchemy 20 | import sqlalchemy.ext.asyncio 21 | 22 | from google.cloud.alloydbconnector import AsyncConnector 23 | 24 | 25 | async def create_sqlalchemy_engine( 26 | inst_uri: str, 27 | user: str, 28 | password: str, 29 | db: str, 30 | ) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, AsyncConnector]: 31 | """Creates a connection pool for an AlloyDB instance and returns the pool 32 | and the connector. Callers are responsible for closing the pool and the 33 | connector. 34 | 35 | A sample invocation looks like: 36 | 37 | engine, connector = await create_sqlalchemy_engine( 38 | inst_uri, 39 | user, 40 | password, 41 | db, 42 | ) 43 | async with engine.connect() as conn: 44 | time = await conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 45 | curr_time = time[0] 46 | # do something with query result 47 | await connector.close() 48 | 49 | Args: 50 | instance_uri (str): 51 | The instance URI specifies the instance relative to the project, 52 | region, and cluster. For example: 53 | "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" 54 | user (str): 55 | The database user name, e.g., postgres 56 | password (str): 57 | The database user's password, e.g., secret-password 58 | db_name (str): 59 | The name of the database, e.g., mydb 60 | """ 61 | connector = AsyncConnector() 62 | 63 | # create SQLAlchemy connection pool 64 | engine = sqlalchemy.ext.asyncio.create_async_engine( 65 | "postgresql+asyncpg://", 66 | async_creator=lambda: connector.connect( 67 | inst_uri, 68 | "asyncpg", 69 | user=user, 70 | password=password, 71 | db=db, 72 | ip_type="PUBLIC", 73 | ), 74 | execution_options={"isolation_level": "AUTOCOMMIT"}, 75 | ) 76 | return engine, connector 77 | 78 | 79 | # [END alloydb_sqlalchemy_connect_async_connector_public_ip] 80 | 81 | 82 | @pytest.mark.asyncio 83 | async def test_connection_with_asyncpg() -> None: 84 | """Basic test to get time from database.""" 85 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 86 | user = os.environ["ALLOYDB_USER"] 87 | password = os.environ["ALLOYDB_PASS"] 88 | db = os.environ["ALLOYDB_DB"] 89 | 90 | pool, connector = await create_sqlalchemy_engine(inst_uri, user, password, db) 91 | 92 | async with pool.connect() as conn: 93 | res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone() 94 | assert res[0] == 1 95 | 96 | await connector.close() 97 | -------------------------------------------------------------------------------- /tests/system/test_native_asyncpg_direct_connection.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # flake8: noqa: ANN001 16 | from datetime import datetime 17 | import os 18 | 19 | # [START alloydb_native_asyncpg_connect_iam_authn_direct] 20 | import asyncpg 21 | import google.auth 22 | from google.auth.transport.requests import Request 23 | 24 | # [END alloydb_native_asyncpg_connect_iam_authn_direct] 25 | 26 | 27 | async def test_native_asyncpg_time() -> None: 28 | """Basic test to get time from database using native asyncpg connection.""" 29 | ip_address = os.environ["ALLOYDB_INSTANCE_IP"] # Private IP for AlloyDB instance 30 | user = os.environ["ALLOYDB_IAM_USER"] 31 | db = os.environ["ALLOYDB_DB"] 32 | 33 | # [START alloydb_native_asyncpg_connect_iam_authn_direct] 34 | # initialize Google Auth credentials 35 | credentials, _ = google.auth.default( 36 | scopes=["https://www.googleapis.com/auth/cloud-platform"] 37 | ) 38 | 39 | def get_authentication_token() -> str: 40 | """Get OAuth2 access token to be used for IAM database authentication""" 41 | # refresh credentials if expired 42 | if not credentials.valid: 43 | request = Request() 44 | credentials.refresh(request) 45 | return credentials.token 46 | 47 | # ... inside of async context (function) 48 | async with asyncpg.create_pool( 49 | user=user, # your IAM db user, e.g. service-account@project-id.iam 50 | password=get_authentication_token, # callable to get fresh OAuth2 token 51 | host=ip_address, # your AlloyDB instance IP address 52 | port=5432, 53 | database=db, # your database name 54 | # Because this connection uses an OAuth2 token as a password, you must 55 | # require SSL, or better, enforce all clients speak SSL on the server 56 | # side. This ensures the OAuth2 token is not inadvertantly leaked. 57 | ssl="require", 58 | ) as pool: 59 | # acquire connection from native asyncpg connection pool 60 | async with pool.acquire() as conn: 61 | time = await conn.fetchrow("SELECT NOW()") 62 | print("Current time is ", time[0]) 63 | # [END alloydb_native_asyncpg_connect_iam_authn_direct] 64 | assert type(time[0]) is datetime 65 | -------------------------------------------------------------------------------- /tests/system/test_pg8000_connection.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from datetime import datetime 16 | import os 17 | 18 | # [START alloydb_sqlalchemy_connect_connector] 19 | import sqlalchemy 20 | 21 | from google.cloud.alloydbconnector import Connector 22 | 23 | 24 | def create_sqlalchemy_engine( 25 | inst_uri: str, 26 | user: str, 27 | password: str, 28 | db: str, 29 | refresh_strategy: str = "background", 30 | ) -> tuple[sqlalchemy.engine.Engine, Connector]: 31 | """Creates a connection pool for an AlloyDB instance and returns the pool 32 | and the connector. Callers are responsible for closing the pool and the 33 | connector. 34 | 35 | A sample invocation looks like: 36 | 37 | engine, connector = create_sqlalchemy_engine( 38 | inst_uri, 39 | user, 40 | password, 41 | db, 42 | ) 43 | with engine.connect() as conn: 44 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 45 | conn.commit() 46 | curr_time = time[0] 47 | # do something with query result 48 | connector.close() 49 | 50 | Args: 51 | instance_uri (str): 52 | The instance URI specifies the instance relative to the project, 53 | region, and cluster. For example: 54 | "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" 55 | user (str): 56 | The database user name, e.g., postgres 57 | password (str): 58 | The database user's password, e.g., secret-password 59 | db (str): 60 | The name of the database, e.g., mydb 61 | refresh_strategy (Optional[str]): 62 | Refresh strategy for the AlloyDB Connector. Can be one of "lazy" 63 | or "background". For serverless environments use "lazy" to avoid 64 | errors resulting from CPU being throttled. 65 | """ 66 | connector = Connector(refresh_strategy=refresh_strategy) 67 | 68 | # create SQLAlchemy connection pool 69 | engine = sqlalchemy.create_engine( 70 | "postgresql+pg8000://", 71 | creator=lambda: connector.connect( 72 | inst_uri, 73 | "pg8000", 74 | user=user, 75 | password=password, 76 | db=db, 77 | ), 78 | ) 79 | engine.dialect.description_encoding = None 80 | return engine, connector 81 | 82 | 83 | # [END alloydb_sqlalchemy_connect_connector] 84 | 85 | 86 | def test_pg8000_connection() -> None: 87 | """Basic test to get time from database.""" 88 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 89 | user = os.environ["ALLOYDB_USER"] 90 | password = os.environ["ALLOYDB_PASS"] 91 | db = os.environ["ALLOYDB_DB"] 92 | 93 | engine, connector = create_sqlalchemy_engine(inst_uri, user, password, db) 94 | with engine.connect() as conn: 95 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 96 | conn.commit() 97 | curr_time = time[0] 98 | assert type(curr_time) is datetime 99 | connector.close() 100 | 101 | 102 | def test_lazy_pg8000_connection() -> None: 103 | """Basic test to get time from database.""" 104 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 105 | user = os.environ["ALLOYDB_USER"] 106 | password = os.environ["ALLOYDB_PASS"] 107 | db = os.environ["ALLOYDB_DB"] 108 | 109 | engine, connector = create_sqlalchemy_engine(inst_uri, user, password, db, "lazy") 110 | with engine.connect() as conn: 111 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 112 | conn.commit() 113 | curr_time = time[0] 114 | assert type(curr_time) is datetime 115 | connector.close() 116 | -------------------------------------------------------------------------------- /tests/system/test_pg8000_iam_authn.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from datetime import datetime 16 | import os 17 | 18 | # [START alloydb_sqlalchemy_connect_connector_iam_authn] 19 | import sqlalchemy 20 | 21 | from google.cloud.alloydbconnector import Connector 22 | 23 | 24 | def create_sqlalchemy_engine( 25 | inst_uri: str, user: str, db: str, refresh_strategy: str = "background" 26 | ) -> tuple[sqlalchemy.engine.Engine, Connector]: 27 | """Creates a connection pool for an AlloyDB instance and returns the pool 28 | and the connector. Callers are responsible for closing the pool and the 29 | connector. 30 | 31 | A sample invocation looks like: 32 | 33 | engine, connector = create_sqlalchemy_engine( 34 | inst_uri, 35 | user, 36 | db, 37 | ) 38 | with engine.connect() as conn: 39 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 40 | conn.commit() 41 | curr_time = time[0] 42 | # do something with query result 43 | connector.close() 44 | 45 | Args: 46 | instance_uri (str): 47 | The instance URI specifies the instance relative to the project, 48 | region, and cluster. For example: 49 | "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" 50 | user (str): 51 | The formatted IAM database username. 52 | e.g., my-email@test.com, service-account@project-id.iam 53 | db (str): 54 | The name of the database, e.g., mydb 55 | refresh_strategy (Optional[str]): 56 | Refresh strategy for the AlloyDB Connector. Can be one of "lazy" 57 | or "background". For serverless environments use "lazy" to avoid 58 | errors resulting from CPU being throttled. 59 | """ 60 | connector = Connector(refresh_strategy=refresh_strategy) 61 | 62 | # create SQLAlchemy connection pool 63 | engine = sqlalchemy.create_engine( 64 | "postgresql+pg8000://", 65 | creator=lambda: connector.connect( 66 | inst_uri, 67 | "pg8000", 68 | user=user, 69 | db=db, 70 | enable_iam_auth=True, 71 | ), 72 | ) 73 | return engine, connector 74 | 75 | 76 | # [END alloydb_sqlalchemy_connect_connector_iam_authn] 77 | 78 | 79 | def test_pg8000_iam_authn_time() -> None: 80 | """Basic test to get time from database.""" 81 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 82 | user = os.environ["ALLOYDB_IAM_USER"] 83 | db = os.environ["ALLOYDB_DB"] 84 | 85 | engine, connector = create_sqlalchemy_engine(inst_uri, user, db) 86 | with engine.connect() as conn: 87 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 88 | conn.commit() 89 | curr_time = time[0] 90 | assert type(curr_time) is datetime 91 | connector.close() 92 | 93 | 94 | def test_pg8000_iam_authn_lazy() -> None: 95 | """Basic test to get time from database.""" 96 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 97 | user = os.environ["ALLOYDB_IAM_USER"] 98 | db = os.environ["ALLOYDB_DB"] 99 | 100 | engine, connector = create_sqlalchemy_engine(inst_uri, user, db, "lazy") 101 | with engine.connect() as conn: 102 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 103 | conn.commit() 104 | curr_time = time[0] 105 | assert type(curr_time) is datetime 106 | connector.close() 107 | -------------------------------------------------------------------------------- /tests/system/test_pg8000_psc.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from datetime import datetime 16 | import os 17 | 18 | import sqlalchemy 19 | 20 | from google.cloud.alloydbconnector import Connector 21 | 22 | 23 | def create_sqlalchemy_engine( 24 | inst_uri: str, 25 | user: str, 26 | password: str, 27 | db: str, 28 | ) -> tuple[sqlalchemy.engine.Engine, Connector]: 29 | """Creates a connection pool for an AlloyDB instance and returns the pool 30 | and the connector. Callers are responsible for closing the pool and the 31 | connector. 32 | 33 | A sample invocation looks like: 34 | 35 | engine, connector = create_sqlalchemy_engine( 36 | inst_uri, 37 | user, 38 | password, 39 | db, 40 | ) 41 | with engine.connect() as conn: 42 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 43 | conn.commit() 44 | curr_time = time[0] 45 | # do something with query result 46 | connector.close() 47 | 48 | Args: 49 | instance_uri (str): 50 | The instance URI specifies the instance relative to the project, 51 | region, and cluster. For example: 52 | "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" 53 | user (str): 54 | The database user name, e.g., postgres 55 | password (str): 56 | The database user's password, e.g., secret-password 57 | db_name (str): 58 | The name of the database, e.g., mydb 59 | """ 60 | connector = Connector() 61 | 62 | # create SQLAlchemy connection pool 63 | engine = sqlalchemy.create_engine( 64 | "postgresql+pg8000://", 65 | creator=lambda: connector.connect( 66 | inst_uri, 67 | "pg8000", 68 | user=user, 69 | password=password, 70 | db=db, 71 | ip_type="PSC", 72 | ), 73 | ) 74 | return engine, connector 75 | 76 | 77 | def test_pg8000_time() -> None: 78 | """Basic test to get time from database.""" 79 | inst_uri = os.environ["ALLOYDB_PSC_INSTANCE_URI"] 80 | user = os.environ["ALLOYDB_USER"] 81 | password = os.environ["ALLOYDB_PASS"] 82 | db = os.environ["ALLOYDB_DB"] 83 | 84 | engine, connector = create_sqlalchemy_engine(inst_uri, user, password, db) 85 | with engine.connect() as conn: 86 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 87 | conn.commit() 88 | curr_time = time[0] 89 | assert type(curr_time) is datetime 90 | connector.close() 91 | -------------------------------------------------------------------------------- /tests/system/test_pg8000_public_ip.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from datetime import datetime 16 | import os 17 | 18 | # [START alloydb_sqlalchemy_connect_connector_public_ip] 19 | import sqlalchemy 20 | 21 | from google.cloud.alloydbconnector import Connector 22 | 23 | 24 | def create_sqlalchemy_engine( 25 | inst_uri: str, 26 | user: str, 27 | password: str, 28 | db: str, 29 | ) -> tuple[sqlalchemy.engine.Engine, Connector]: 30 | """Creates a connection pool for an AlloyDB instance and returns the pool 31 | and the connector. Callers are responsible for closing the pool and the 32 | connector. 33 | 34 | A sample invocation looks like: 35 | 36 | engine, connector = create_sqlalchemy_engine( 37 | inst_uri, 38 | user, 39 | password, 40 | db, 41 | ) 42 | with engine.connect() as conn: 43 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 44 | conn.commit() 45 | curr_time = time[0] 46 | # do something with query result 47 | connector.close() 48 | 49 | Args: 50 | instance_uri (str): 51 | The instance URI specifies the instance relative to the project, 52 | region, and cluster. For example: 53 | "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" 54 | user (str): 55 | The database user name, e.g., postgres 56 | password (str): 57 | The database user's password, e.g., secret-password 58 | db_name (str): 59 | The name of the database, e.g., mydb 60 | """ 61 | connector = Connector() 62 | 63 | # create SQLAlchemy connection pool 64 | engine = sqlalchemy.create_engine( 65 | "postgresql+pg8000://", 66 | creator=lambda: connector.connect( 67 | inst_uri, 68 | "pg8000", 69 | user=user, 70 | password=password, 71 | db=db, 72 | ip_type="PUBLIC", 73 | ), 74 | ) 75 | return engine, connector 76 | 77 | 78 | # [END alloydb_sqlalchemy_connect_connector_public_ip] 79 | 80 | 81 | def test_pg8000_time() -> None: 82 | """Basic test to get time from database.""" 83 | inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] 84 | user = os.environ["ALLOYDB_USER"] 85 | password = os.environ["ALLOYDB_PASS"] 86 | db = os.environ["ALLOYDB_DB"] 87 | 88 | engine, connector = create_sqlalchemy_engine(inst_uri, user, password, db) 89 | with engine.connect() as conn: 90 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 91 | conn.commit() 92 | curr_time = time[0] 93 | assert type(curr_time) is datetime 94 | connector.close() 95 | -------------------------------------------------------------------------------- /tests/system/test_psycopg2_direct_connection.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # flake8: noqa: ANN001 16 | from datetime import datetime 17 | import os 18 | 19 | # [START alloydb_psycopg2_connect_iam_authn_direct] 20 | import google.auth 21 | from google.auth.credentials import Credentials 22 | from google.auth.transport.requests import Request 23 | import sqlalchemy 24 | from sqlalchemy import event 25 | 26 | # [END alloydb_psycopg2_connect_iam_authn_direct] 27 | 28 | 29 | def create_sqlalchemy_engine( 30 | ip_address: str, 31 | user: str, 32 | db_name: str, 33 | ) -> sqlalchemy.engine.Engine: 34 | """Creates a SQLAlchemy connection pool for an AlloyDB instance configured 35 | using psycopg2. 36 | 37 | Callers are responsible for closing the pool. This implementation uses a 38 | direct TCP connection with IAM database authentication and not 39 | the Cloud SQL Python Connector. 40 | 41 | A sample invocation looks like: 42 | 43 | engine = create_sqlalchemy_engine( 44 | ip_address, 45 | user, 46 | db, 47 | ) 48 | 49 | with engine.connect() as conn: 50 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 51 | conn.commit() 52 | 53 | Args: 54 | ip_address (str): 55 | The IP address of an AlloyDB instance, e.g., 10.0.0.1 56 | user (str): 57 | The formatted IAM database username. 58 | e.g., my-email@test.com, service-account@project-id.iam 59 | db_name (str): 60 | The name of the database, e.g., mydb 61 | """ 62 | # [START alloydb_psycopg2_connect_iam_authn_direct] 63 | # initialize Google Auth creds 64 | creds, _ = google.auth.default( 65 | scopes=["https://www.googleapis.com/auth/cloud-platform"] 66 | ) 67 | 68 | def get_authentication_token(credentials: Credentials) -> str: 69 | """Get OAuth2 access token to be used for IAM database authentication""" 70 | # refresh credentials if expired 71 | if not credentials.valid: 72 | request = Request() 73 | credentials.refresh(request) 74 | return credentials.token 75 | 76 | engine = sqlalchemy.create_engine( 77 | # Equivalent URL: 78 | # postgresql+psycopg2://:empty@:5432/ 79 | sqlalchemy.engine.url.URL.create( 80 | drivername="postgresql+psycopg2", 81 | username=user, # IAM db user, e.g. service-account@project-id.iam 82 | password="", # placeholder to be replaced with OAuth2 token 83 | host=ip_address, # AlloyDB instance IP address 84 | port=5432, 85 | database=db_name, # "my-database-name" 86 | ), 87 | connect_args={"sslmode": "require"}, 88 | ) 89 | 90 | # set 'do_connect' event listener to replace password with OAuth2 token 91 | @event.listens_for(engine, "do_connect") 92 | def auto_iam_authentication(dialect, conn_rec, cargs, cparams) -> None: 93 | cparams["password"] = get_authentication_token(creds) 94 | 95 | # [END alloydb_psycopg2_connect_iam_authn_direct] 96 | return engine 97 | 98 | 99 | def test_psycopg2_time() -> None: 100 | """Basic test to get time from database.""" 101 | ip_address = os.environ["ALLOYDB_INSTANCE_IP"] # Private IP for AlloyDB instance 102 | user = os.environ["ALLOYDB_IAM_USER"] 103 | db = os.environ["ALLOYDB_DB"] 104 | 105 | engine = create_sqlalchemy_engine(ip_address, user, db) 106 | # [START alloydb_psycopg2_connect_iam_authn_direct] 107 | # use connection from connection pool to query AlloyDB database 108 | with engine.connect() as conn: 109 | time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() 110 | conn.commit() 111 | print("Current time is ", time[0]) 112 | # [END alloydb_psycopg2_connect_iam_authn_direct] 113 | curr_time = time[0] 114 | assert type(curr_time) is datetime 115 | -------------------------------------------------------------------------------- /tests/system/test_sqlalchemy_asyncpg_direct_connection.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # flake8: noqa: ANN001 16 | from datetime import datetime 17 | import os 18 | 19 | # [START alloydb_sqlalchemy_asyncpg_connect_iam_authn_direct] 20 | import google.auth 21 | from google.auth.credentials import Credentials 22 | from google.auth.transport.requests import Request 23 | import sqlalchemy 24 | from sqlalchemy import event 25 | from sqlalchemy.ext.asyncio import create_async_engine 26 | 27 | # [END alloydb_sqlalchemy_asyncpg_connect_iam_authn_direct] 28 | 29 | 30 | def create_sqlalchemy_engine( 31 | ip_address: str, 32 | user: str, 33 | db_name: str, 34 | ) -> sqlalchemy.ext.asyncio.engine.AsyncEngine: 35 | """Creates a SQLAlchemy connection pool for an AlloyDB instance configured 36 | using asyncpg. 37 | 38 | Callers are responsible for closing the pool. This implementation uses a 39 | direct TCP connection with IAM database authentication and not 40 | the Cloud SQL Python Connector. 41 | 42 | A sample invocation looks like: 43 | 44 | engine = create_sqlalchemy_engine( 45 | ip_address, 46 | user, 47 | db, 48 | ) 49 | 50 | async with engine.connect() as conn: 51 | result = await conn.execute(sqlalchemy.text("SELECT NOW()")) 52 | conn.commit() 53 | time = result.fetchone() 54 | 55 | Args: 56 | ip_address (str): 57 | The IP address of an AlloyDB instance, e.g., 10.0.0.1 58 | user (str): 59 | The formatted IAM database username. 60 | e.g., my-email@test.com, service-account@project-id.iam 61 | db_name (str): 62 | The name of the database, e.g., mydb 63 | """ 64 | # [START alloydb_sqlalchemy_asyncpg_connect_iam_authn_direct] 65 | # initialize Google Auth credentials 66 | credentials, _ = google.auth.default( 67 | scopes=["https://www.googleapis.com/auth/cloud-platform"] 68 | ) 69 | 70 | def get_authentication_token(credentials: Credentials) -> str: 71 | """Get OAuth2 access token to be used for IAM database authentication""" 72 | # refresh credentials if expired 73 | if not credentials.valid: 74 | request = Request() 75 | credentials.refresh(request) 76 | return credentials.token 77 | 78 | engine = create_async_engine( 79 | # Equivalent URL: 80 | # postgresql+asyncpg://:empty@:5432/ 81 | sqlalchemy.engine.url.URL.create( 82 | drivername="postgresql+asyncpg", 83 | username=user, # your IAM db user, e.g. service-account@project-id.iam 84 | password="", # placeholder to be replaced with OAuth2 token 85 | host=ip_address, # your AlloyDB instance IP address 86 | port=5432, 87 | database=db_name, # your database name 88 | ), 89 | # Because this connection uses an OAuth2 token as a password, you must 90 | # require SSL, or better, enforce all clients speak SSL on the server 91 | # side. This ensures the OAuth2 token is not inadvertantly leaked. 92 | connect_args={"ssl": "require"}, 93 | ) 94 | 95 | # set 'do_connect' event listener to replace password with OAuth2 token 96 | # must use engine.sync_engine as async events are not implemented 97 | @event.listens_for(engine.sync_engine, "do_connect") 98 | def auto_iam_authentication(dialect, conn_rec, cargs, cparams) -> None: 99 | cparams["password"] = get_authentication_token(credentials) 100 | 101 | # [END alloydb_sqlalchemy_asyncpg_connect_iam_authn_direct] 102 | return engine 103 | 104 | 105 | async def test_sqlalchemy_asyncpg_time() -> None: 106 | """Basic test to get time from database using asyncpg with SQLAlchemy.""" 107 | ip_address = os.environ["ALLOYDB_INSTANCE_IP"] # Private IP for AlloyDB instance 108 | user = os.environ["ALLOYDB_IAM_USER"] 109 | db = os.environ["ALLOYDB_DB"] 110 | 111 | engine = create_sqlalchemy_engine(ip_address, user, db) 112 | # [START alloydb_sqlalchemy_asyncpg_connect_iam_authn_direct] 113 | # use connection from connection pool to query AlloyDB database 114 | async with engine.connect() as conn: 115 | result = await conn.execute(sqlalchemy.text("SELECT NOW()")) 116 | time = result.fetchone() 117 | print("Current time is ", time[0]) 118 | # [END alloydb_sqlalchemy_asyncpg_connect_iam_authn_direct] 119 | curr_time = time[0] 120 | assert type(curr_time) is datetime 121 | # cleanup AsyncEngine 122 | await engine.dispose() 123 | -------------------------------------------------------------------------------- /tests/unit/conftest.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import asyncio 16 | import socket 17 | import ssl 18 | from threading import Thread 19 | 20 | from aiofiles.tempfile import TemporaryDirectory 21 | from mocks import FakeAlloyDBClient 22 | from mocks import FakeCredentials 23 | from mocks import FakeInstance 24 | from mocks import metadata_exchange 25 | import pytest 26 | 27 | from google.cloud.alloydbconnector.utils import _write_to_file 28 | 29 | DELAY = 1.0 30 | 31 | 32 | @pytest.fixture 33 | def credentials() -> FakeCredentials: 34 | return FakeCredentials() 35 | 36 | 37 | @pytest.fixture(scope="session") 38 | def fake_instance() -> FakeInstance: 39 | return FakeInstance() 40 | 41 | 42 | @pytest.fixture 43 | def fake_client(fake_instance: FakeInstance) -> FakeAlloyDBClient: 44 | return FakeAlloyDBClient(fake_instance) 45 | 46 | 47 | async def start_proxy_server(instance: FakeInstance) -> None: 48 | """Run local proxy server capable of performing metadata exchange""" 49 | ip_address = "127.0.0.1" 50 | port = 5433 51 | # create socket 52 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: 53 | # create SSL/TLS context 54 | context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) 55 | context.minimum_version = ssl.TLSVersion.TLSv1_3 56 | root, _, server = instance.get_pem_certs() 57 | # tmpdir and its contents are automatically deleted after the CA cert 58 | # and cert chain are loaded into the SSLcontext. The values 59 | # need to be written to files in order to be loaded by the SSLContext 60 | async with TemporaryDirectory() as tmpdir: 61 | _, cert_chain_filename, key_filename = await _write_to_file( 62 | tmpdir, server, [server, root], instance.server_key 63 | ) 64 | context.load_cert_chain(cert_chain_filename, key_filename) 65 | # bind socket to AlloyDB proxy server port on localhost 66 | sock.bind((ip_address, port)) 67 | # listen for incoming connections 68 | sock.listen(5) 69 | 70 | with context.wrap_socket(sock, server_side=True) as ssock: 71 | while True: 72 | conn, _ = ssock.accept() 73 | metadata_exchange(conn) 74 | conn.sendall(instance.name.encode("utf-8")) 75 | conn.close() 76 | 77 | 78 | @pytest.fixture(scope="session") 79 | def proxy_server(fake_instance: FakeInstance) -> None: 80 | """Run local proxy server capable of performing metadata exchange""" 81 | thread = Thread( 82 | target=asyncio.run, 83 | args=( 84 | start_proxy_server( 85 | fake_instance, 86 | ), 87 | ), 88 | daemon=True, 89 | ) 90 | thread.start() 91 | thread.join(DELAY) # add a delay to allow the proxy server to start 92 | -------------------------------------------------------------------------------- /tests/unit/test_client.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from typing import Optional 16 | 17 | import google.cloud.alloydb_v1beta as v1beta 18 | from mocks import FakeAlloyDBAdminAsyncClient 19 | from mocks import FakeAlloyDBAdminClient 20 | from mocks import FakeCredentials 21 | import pytest 22 | 23 | from google.cloud.alloydbconnector.client import AlloyDBClient 24 | from google.cloud.alloydbconnector.utils import generate_keys 25 | from google.cloud.alloydbconnector.version import __version__ as version 26 | 27 | 28 | @pytest.mark.asyncio 29 | async def test__get_metadata(credentials: FakeCredentials) -> None: 30 | """ 31 | Test _get_metadata returns successfully. 32 | """ 33 | test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminAsyncClient()) 34 | ip_addrs = await test_client._get_metadata( 35 | "test-project", 36 | "test-region", 37 | "test-cluster", 38 | "test-instance", 39 | ) 40 | assert ip_addrs == { 41 | "PRIVATE": "10.0.0.1", 42 | "PUBLIC": "", 43 | "PSC": "", 44 | } 45 | 46 | 47 | @pytest.mark.asyncio 48 | async def test__get_metadata_with_public_ip(credentials: FakeCredentials) -> None: 49 | """ 50 | Test _get_metadata returns successfully with Public IP. 51 | """ 52 | test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminAsyncClient()) 53 | ip_addrs = await test_client._get_metadata( 54 | "test-project", 55 | "test-region", 56 | "test-cluster", 57 | "public-instance", 58 | ) 59 | assert ip_addrs == { 60 | "PRIVATE": "10.0.0.1", 61 | "PUBLIC": "127.0.0.1", 62 | "PSC": "", 63 | } 64 | 65 | 66 | @pytest.mark.asyncio 67 | async def test__get_metadata_with_psc(credentials: FakeCredentials) -> None: 68 | """ 69 | Test _get_metadata returns successfully with PSC DNS name. 70 | """ 71 | test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminAsyncClient()) 72 | ip_addrs = await test_client._get_metadata( 73 | "test-project", 74 | "test-region", 75 | "test-cluster", 76 | "psc-instance", 77 | ) 78 | assert ip_addrs == { 79 | "PRIVATE": "", 80 | "PUBLIC": "", 81 | "PSC": "x.y.alloydb.goog", 82 | } 83 | 84 | 85 | async def test__get_metadata_with_async_client(credentials: FakeCredentials) -> None: 86 | """ 87 | Test _get_metadata returns successfully for an async client. 88 | """ 89 | test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminAsyncClient()) 90 | test_client._is_sync = False 91 | assert ( 92 | await test_client._get_metadata( 93 | "test-project", 94 | "test-region", 95 | "test-cluster", 96 | "psc-instance", 97 | ) 98 | is not None 99 | ) 100 | 101 | 102 | async def test__get_metadata_with_sync_client(credentials: FakeCredentials) -> None: 103 | """ 104 | Test _get_metadata returns successfully for a sync client. 105 | """ 106 | test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminClient()) 107 | test_client._is_sync = True 108 | assert ( 109 | await test_client._get_metadata( 110 | "test-project", 111 | "test-region", 112 | "test-cluster", 113 | "psc-instance", 114 | ) 115 | is not None 116 | ) 117 | 118 | 119 | @pytest.mark.asyncio 120 | async def test__get_client_certificate(credentials: FakeCredentials) -> None: 121 | """ 122 | Test _get_client_certificate returns successfully. 123 | """ 124 | test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminAsyncClient()) 125 | keys = await generate_keys() 126 | certs = await test_client._get_client_certificate( 127 | "test-project", "test-region", "test-cluster", keys[1] 128 | ) 129 | ca_cert, cert_chain = certs 130 | assert ca_cert == "This is the CA cert" 131 | assert cert_chain[0] == "This is the client cert" 132 | assert cert_chain[1] == "This is the intermediate cert" 133 | assert cert_chain[2] == "This is the root cert" 134 | 135 | 136 | async def test__get_client_certificate_with_async_client( 137 | credentials: FakeCredentials, 138 | ) -> None: 139 | """ 140 | Test _get_client_certificate returns successfully for an async client. 141 | """ 142 | test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminAsyncClient()) 143 | test_client._is_sync = False 144 | keys = await generate_keys() 145 | assert ( 146 | await test_client._get_client_certificate( 147 | "test-project", "test-region", "test-cluster", keys[1] 148 | ) 149 | is not None 150 | ) 151 | 152 | 153 | async def test__get_client_certificate_with_sync_client( 154 | credentials: FakeCredentials, 155 | ) -> None: 156 | """ 157 | Test _get_client_certificate returns successfully for a sync client. 158 | """ 159 | test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminClient()) 160 | test_client._is_sync = True 161 | keys = await generate_keys() 162 | assert ( 163 | await test_client._get_client_certificate( 164 | "test-project", "test-region", "test-cluster", keys[1] 165 | ) 166 | is not None 167 | ) 168 | 169 | 170 | @pytest.mark.asyncio 171 | async def test_AlloyDBClient_init_(credentials: FakeCredentials) -> None: 172 | """ 173 | Test to check whether the __init__ method of AlloyDBClient 174 | can correctly initialize a client. 175 | """ 176 | client = AlloyDBClient("www.test-endpoint.com", "my-quota-project", credentials) 177 | # verify base endpoint is set 178 | assert client._client.api_endpoint == "www.test-endpoint.com" 179 | # verify proper headers are set 180 | assert client._user_agent.startswith(f"alloydb-python-connector/{version}") 181 | assert client._client._client._client_options.quota_project_id == "my-quota-project" 182 | 183 | 184 | @pytest.mark.asyncio 185 | async def test_AlloyDBClient_init_custom_user_agent( 186 | credentials: FakeCredentials, 187 | ) -> None: 188 | """ 189 | Test to check that custom user agents are included in HTTP requests. 190 | """ 191 | client = AlloyDBClient( 192 | "www.test-endpoint.com", 193 | "my-quota-project", 194 | credentials, 195 | user_agent="custom-agent/v1.0.0 other-agent/v2.0.0", 196 | ) 197 | assert client._user_agent.startswith( 198 | f"alloydb-python-connector/{version} custom-agent/v1.0.0 other-agent/v2.0.0" 199 | ) 200 | 201 | 202 | async def test_AlloyDBClient_init_specified_client( 203 | credentials: FakeCredentials, 204 | ) -> None: 205 | """ 206 | Test to check that __init__ method of AlloyDBClient uses specified client. 207 | """ 208 | client = AlloyDBClient( 209 | "www.test-endpoint.com", 210 | "my-quota-project", 211 | credentials, 212 | FakeAlloyDBAdminAsyncClient(), 213 | ) 214 | assert client._is_sync is False 215 | assert type(client._client) is FakeAlloyDBAdminAsyncClient 216 | 217 | 218 | async def test_AlloyDBClient_init_sync_client(credentials: FakeCredentials) -> None: 219 | """ 220 | Test to check that __init__ method of AlloyDBClient creates a sync client 221 | when client is not specified and driver is pg8000. 222 | """ 223 | client = AlloyDBClient( 224 | "www.test-endpoint.com", "my-quota-project", credentials, driver="pg8000" 225 | ) 226 | assert client._is_sync is True 227 | assert type(client._client) is v1beta.AlloyDBAdminClient 228 | assert client._client.transport.kind == "grpc" 229 | 230 | 231 | async def test_AlloyDBClient_init_async_client(credentials: FakeCredentials) -> None: 232 | """ 233 | Test to check that __init__ method of AlloyDBClient creates an async client 234 | when client is not specified and driver is not pg8000. 235 | """ 236 | client = AlloyDBClient( 237 | "www.test-endpoint.com", "my-quota-project", credentials, driver="" 238 | ) 239 | assert client._is_sync is False 240 | assert type(client._client) is v1beta.AlloyDBAdminAsyncClient 241 | assert client._client.transport.kind == "grpc_asyncio" 242 | 243 | 244 | @pytest.mark.parametrize( 245 | "driver", 246 | [None, "pg8000", "asyncpg"], 247 | ) 248 | @pytest.mark.asyncio 249 | async def test_AlloyDBClient_user_agent( 250 | driver: Optional[str], credentials: FakeCredentials 251 | ) -> None: 252 | """ 253 | Test to check whether the __init__ method of AlloyDBClient 254 | properly sets user agent when passed a database driver. 255 | """ 256 | client = AlloyDBClient( 257 | "www.test-endpoint.com", "my-quota-project", credentials, driver=driver 258 | ) 259 | if driver is None: 260 | assert client._user_agent.startswith(f"alloydb-python-connector/{version}") 261 | else: 262 | assert client._user_agent.startswith( 263 | f"alloydb-python-connector/{version}+{driver}" 264 | ) 265 | 266 | 267 | @pytest.mark.parametrize( 268 | "driver, expected", 269 | [(None, False), ("pg8000", True), ("asyncpg", False)], 270 | ) 271 | @pytest.mark.asyncio 272 | async def test_AlloyDBClient_use_metadata( 273 | driver: Optional[str], expected: bool, credentials: FakeCredentials 274 | ) -> None: 275 | """ 276 | Test to check whether the __init__ method of AlloyDBClient 277 | properly sets use_metadata. 278 | """ 279 | client = AlloyDBClient( 280 | "www.test-endpoint.com", "my-quota-project", credentials, driver=driver 281 | ) 282 | assert client._use_metadata == expected 283 | -------------------------------------------------------------------------------- /tests/unit/test_connection_info.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from datetime import datetime 16 | from datetime import timedelta 17 | from datetime import timezone 18 | import ssl 19 | 20 | from cryptography import x509 21 | from cryptography.hazmat.primitives import hashes 22 | from cryptography.hazmat.primitives import serialization 23 | from cryptography.hazmat.primitives.asymmetric import rsa 24 | from mocks import FakeInstance 25 | import pytest 26 | 27 | from google.cloud.alloydbconnector import IPTypes 28 | from google.cloud.alloydbconnector.connection_info import ConnectionInfo 29 | from google.cloud.alloydbconnector.exceptions import IPTypeNotFoundError 30 | 31 | 32 | async def test_ConnectionInfo_init_(fake_instance: FakeInstance) -> None: 33 | """ 34 | Test to check whether the __init__ method of ConnectionInfo 35 | can correctly initialize TLS context. 36 | """ 37 | key = rsa.generate_private_key(public_exponent=65537, key_size=2048) 38 | root_cert, intermediate_cert, ca_cert = fake_instance.get_pem_certs() 39 | # build client cert 40 | client_cert = ( 41 | x509.CertificateBuilder() 42 | .subject_name(fake_instance.intermediate_cert.subject) 43 | .issuer_name(fake_instance.intermediate_cert.issuer) 44 | .public_key(key.public_key()) 45 | .serial_number(x509.random_serial_number()) 46 | .not_valid_before(datetime.now(timezone.utc)) 47 | .not_valid_after(datetime.now(timezone.utc) + timedelta(minutes=10)) 48 | ) 49 | # sign client cert with intermediate cert 50 | client_cert = client_cert.sign(fake_instance.intermediate_key, hashes.SHA256()) 51 | client_cert = client_cert.public_bytes(encoding=serialization.Encoding.PEM).decode( 52 | "UTF-8" 53 | ) 54 | conn_info = ConnectionInfo( 55 | [client_cert, intermediate_cert, root_cert], 56 | ca_cert, 57 | key, 58 | fake_instance.ip_addrs, 59 | datetime.now(timezone.utc) + timedelta(minutes=10), 60 | ) 61 | context = await conn_info.create_ssl_context() 62 | # verify TLS requirements 63 | assert context.minimum_version == ssl.TLSVersion.TLSv1_3 64 | 65 | 66 | async def test_ConnectionInfo_caches_sslcontext() -> None: 67 | info = ConnectionInfo(["cert"], "cert", "key".encode(), {}, datetime.now()) 68 | # context should default to None 69 | assert info.context is None 70 | # cache a 'context' 71 | info.context = "context" 72 | # calling create_ssl_context should no-op with an existing 'context' 73 | await info.create_ssl_context() 74 | assert info.context == "context" 75 | 76 | 77 | @pytest.mark.parametrize( 78 | "ip_type, expected", 79 | [ 80 | ( 81 | IPTypes.PRIVATE, 82 | "127.0.0.1", 83 | ), 84 | ( 85 | IPTypes.PUBLIC, 86 | "0.0.0.0", 87 | ), 88 | ( 89 | IPTypes.PSC, 90 | "x.y.alloydb.goog", 91 | ), 92 | ], 93 | ) 94 | async def test_ConnectionInfo_get_preferred_ip(ip_type: IPTypes, expected: str) -> None: 95 | """Test that ConnectionInfo.get_preferred_ip returns proper ip address.""" 96 | ip_addrs = { 97 | "PRIVATE": "127.0.0.1", 98 | "PUBLIC": "0.0.0.0", 99 | "PSC": "x.y.alloydb.goog", 100 | } 101 | conn_info = ConnectionInfo( 102 | ["cert"], "cert", "key", ip_addrs, datetime.now(timezone.utc) 103 | ) 104 | ip_address = conn_info.get_preferred_ip(ip_type) 105 | assert ip_address == expected 106 | 107 | 108 | async def test_ConnectionInfo_get_preferred_ip_IPTypeNotFoundError() -> None: 109 | """Test that ConnectionInfo.get_preferred_ip throws IPTypeNotFoundError""" 110 | conn_info = ConnectionInfo( 111 | ["cert"], 112 | "cert", 113 | "key", 114 | {}, 115 | datetime.now(timezone.utc), 116 | ) 117 | # check error is thrown 118 | with pytest.raises(IPTypeNotFoundError): 119 | conn_info.get_preferred_ip(ip_type=IPTypes.PUBLIC) 120 | -------------------------------------------------------------------------------- /tests/unit/test_instance.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import asyncio 16 | from datetime import datetime 17 | from datetime import timedelta 18 | 19 | import aiohttp 20 | from mocks import FakeAlloyDBClient 21 | import pytest 22 | 23 | from google.cloud.alloydbconnector.connection_info import ConnectionInfo 24 | from google.cloud.alloydbconnector.exceptions import RefreshError 25 | from google.cloud.alloydbconnector.instance import _parse_instance_uri 26 | from google.cloud.alloydbconnector.instance import RefreshAheadCache 27 | from google.cloud.alloydbconnector.refresh_utils import _is_valid 28 | from google.cloud.alloydbconnector.utils import generate_keys 29 | 30 | 31 | @pytest.mark.parametrize( 32 | "instance_uri, expected", 33 | [ 34 | ( 35 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 36 | ("test-project", "test-region", "test-cluster", "test-instance"), 37 | ), 38 | ( 39 | "projects/test-domain:test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 40 | ( 41 | "test-domain:test-project", 42 | "test-region", 43 | "test-cluster", 44 | "test-instance", 45 | ), 46 | ), 47 | ], 48 | ) 49 | def test_parse_instance_uri( 50 | instance_uri: str, expected: tuple[str, str, str, str] 51 | ) -> None: 52 | """ 53 | Test that _parse_instance_uri works correctly on 54 | normal instance uri and domain-scoped projects. 55 | """ 56 | assert expected == _parse_instance_uri(instance_uri) 57 | 58 | 59 | def test_parse_bad_instance_uri() -> None: 60 | """ 61 | Tests that ValueError is thrown for bad instance uri. 62 | """ 63 | with pytest.raises(ValueError): 64 | _parse_instance_uri("test-project:test-instance") 65 | 66 | 67 | @pytest.mark.asyncio 68 | async def test_RefreshAheadCache_init() -> None: 69 | """ 70 | Test to check whether the __init__ method of RefreshAheadCache 71 | can tell if the instance URI that's passed in is formatted correctly. 72 | """ 73 | keys = asyncio.create_task(generate_keys()) 74 | async with aiohttp.ClientSession() as client: 75 | cache = RefreshAheadCache( 76 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 77 | client, 78 | keys, 79 | ) 80 | assert ( 81 | cache._project == "test-project" 82 | and cache._region == "test-region" 83 | and cache._cluster == "test-cluster" 84 | and cache._name == "test-instance" 85 | ) 86 | 87 | 88 | @pytest.mark.asyncio 89 | async def test_RefreshAheadCache_init_invalid_instant_uri() -> None: 90 | """ 91 | Test to check whether the __init__ method of RefreshAheadCache 92 | will throw error for invalid instance URI. 93 | """ 94 | keys = asyncio.create_task(generate_keys()) 95 | async with aiohttp.ClientSession() as client: 96 | with pytest.raises(ValueError): 97 | RefreshAheadCache("invalid/instance/uri/", client, keys) 98 | 99 | 100 | @pytest.mark.asyncio 101 | async def test_RefreshAheadCache_close() -> None: 102 | """ 103 | Test that RefreshAheadCache's close method 104 | cancels tasks gracefully. 105 | """ 106 | keys = asyncio.create_task(generate_keys()) 107 | client = FakeAlloyDBClient() 108 | cache = RefreshAheadCache( 109 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 110 | client, 111 | keys, 112 | ) 113 | # make sure tasks aren't cancelled 114 | assert cache._current.cancelled() is False 115 | assert cache._next.cancelled() is False 116 | # run close() to cancel tasks 117 | await cache.close() 118 | # verify tasks are cancelled 119 | assert (cache._current.done() or cache._current.cancelled()) is True 120 | assert cache._next.cancelled() is True 121 | 122 | 123 | @pytest.mark.asyncio 124 | async def test_perform_refresh() -> None: 125 | """Test that _perform refresh returns valid ConnectionInfo""" 126 | keys = asyncio.create_task(generate_keys()) 127 | client = FakeAlloyDBClient() 128 | cache = RefreshAheadCache( 129 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 130 | client, 131 | keys, 132 | ) 133 | refresh = await cache._perform_refresh() 134 | assert refresh.ip_addrs == { 135 | "PRIVATE": "127.0.0.1", 136 | "PUBLIC": "0.0.0.0", 137 | "PSC": "x.y.alloydb.goog", 138 | } 139 | assert refresh.expiration == client.instance.cert_expiry.replace(microsecond=0) 140 | # close instance 141 | await cache.close() 142 | 143 | 144 | @pytest.mark.asyncio 145 | async def test_schedule_refresh_replaces_result() -> None: 146 | """ 147 | Test to check whether _schedule_refresh replaces a valid refresh result 148 | with another refresh result. 149 | """ 150 | keys = asyncio.create_task(generate_keys()) 151 | client = FakeAlloyDBClient() 152 | cache = RefreshAheadCache( 153 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 154 | client, 155 | keys, 156 | ) 157 | # check current refresh is valid 158 | assert await _is_valid(cache._current) is True 159 | current_refresh = cache._current 160 | # schedule new refresh 161 | await cache._schedule_refresh(0) 162 | new_refresh = cache._current 163 | # verify current has been replaced with new refresh 164 | assert current_refresh != new_refresh 165 | # check new refresh is valid 166 | assert await _is_valid(new_refresh) is True 167 | # close instance 168 | await cache.close() 169 | 170 | 171 | @pytest.mark.asyncio 172 | async def test_schedule_refresh_wont_replace_valid_result_with_invalid() -> None: 173 | """ 174 | Test to check whether _schedule_refresh won't replace a valid 175 | refresh result with an invalid one. 176 | """ 177 | keys = asyncio.create_task(generate_keys()) 178 | client = FakeAlloyDBClient() 179 | cache = RefreshAheadCache( 180 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 181 | client, 182 | keys, 183 | ) 184 | # check current refresh is valid 185 | assert await _is_valid(cache._current) is True 186 | current_refresh = cache._current 187 | # set certificate to be expired 188 | client.instance.cert_before = datetime.now() - timedelta(minutes=20) 189 | client.instance.cert_expiry = datetime.now() - timedelta(minutes=10) 190 | # schedule new refresh 191 | new_refresh = cache._schedule_refresh(0) 192 | # check new refresh is invalid 193 | assert await _is_valid(new_refresh) is False 194 | # check current was not replaced 195 | assert current_refresh == cache._current 196 | # close instance 197 | await cache.close() 198 | 199 | 200 | @pytest.mark.asyncio 201 | async def test_schedule_refresh_expired_cert() -> None: 202 | """ 203 | Test to check whether _schedule_refresh will throw RefreshError on 204 | expired certificate. 205 | """ 206 | keys = asyncio.create_task(generate_keys()) 207 | client = FakeAlloyDBClient() 208 | # set certificate to be expired 209 | client.instance.cert_before = datetime.now() - timedelta(minutes=20) 210 | client.instance.cert_expiry = datetime.now() - timedelta(minutes=10) 211 | cache = RefreshAheadCache( 212 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 213 | client, 214 | keys, 215 | ) 216 | # check RefreshError is thrown 217 | with pytest.raises(RefreshError): 218 | await cache._current 219 | # close instance 220 | await cache.close() 221 | 222 | 223 | @pytest.mark.asyncio 224 | async def test_force_refresh_cancels_pending_refresh() -> None: 225 | """ 226 | Test that force_refresh cancels pending task if refresh_in_progress event is not set. 227 | """ 228 | keys = asyncio.create_task(generate_keys()) 229 | client = FakeAlloyDBClient() 230 | cache = RefreshAheadCache( 231 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 232 | client, 233 | keys, 234 | ) 235 | # make sure initial refresh is finished 236 | await cache._current 237 | # since the pending refresh isn't for another ~56 min, the refresh_in_progress event 238 | # shouldn't be set 239 | pending_refresh = cache._next 240 | assert cache._refresh_in_progress.is_set() is False 241 | await cache.force_refresh() 242 | # pending_refresh has to be awaited for it to raised as cancelled 243 | with pytest.raises(asyncio.CancelledError): 244 | assert await pending_refresh 245 | # verify pending_refresh has now been cancelled 246 | assert pending_refresh.cancelled() is True 247 | assert isinstance(await cache._current, ConnectionInfo) 248 | # close instance 249 | await cache.close() 250 | -------------------------------------------------------------------------------- /tests/unit/test_lazy.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import asyncio 16 | 17 | from google.cloud.alloydbconnector.client import AlloyDBClient 18 | from google.cloud.alloydbconnector.connection_info import ConnectionInfo 19 | from google.cloud.alloydbconnector.lazy import LazyRefreshCache 20 | from google.cloud.alloydbconnector.utils import generate_keys 21 | 22 | 23 | async def test_LazyRefreshCache_connect_info(fake_client: AlloyDBClient) -> None: 24 | """ 25 | Test that LazyRefreshCache.connect_info works as expected. 26 | """ 27 | keys = asyncio.create_task(generate_keys()) 28 | cache = LazyRefreshCache( 29 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 30 | client=fake_client, 31 | keys=keys, 32 | ) 33 | # check that cached connection info is empty 34 | assert cache._cached is None 35 | conn_info = await cache.connect_info() 36 | # check that cached connection info is now set 37 | assert isinstance(cache._cached, ConnectionInfo) 38 | # check that calling connect_info uses cached info 39 | conn_info2 = await cache.connect_info() 40 | assert conn_info2 == conn_info 41 | 42 | 43 | async def test_LazyRefreshCache_force_refresh(fake_client: AlloyDBClient) -> None: 44 | """ 45 | Test that LazyRefreshCache.force_refresh works as expected. 46 | """ 47 | keys = asyncio.create_task(generate_keys()) 48 | cache = LazyRefreshCache( 49 | "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", 50 | client=fake_client, 51 | keys=keys, 52 | ) 53 | conn_info = await cache.connect_info() 54 | # check that cached connection info is now set 55 | assert isinstance(cache._cached, ConnectionInfo) 56 | await cache.force_refresh() 57 | # check that calling connect_info after force_refresh gets new ConnectionInfo 58 | conn_info2 = await cache.connect_info() 59 | # check that new connection info was retrieved 60 | assert conn_info2 != conn_info 61 | assert cache._cached == conn_info2 62 | await cache.close() 63 | -------------------------------------------------------------------------------- /tests/unit/test_packaging.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | import pathlib 17 | import subprocess 18 | import sys 19 | 20 | 21 | def test_namespace_package_compat(tmp_path: pathlib.PosixPath) -> None: 22 | # The ``google`` namespace package should not be masked by the presence of 23 | # `google-cloud-alloydb` and ``google-cloud-alloydb-connector``. 24 | google = tmp_path / "google" 25 | google.mkdir() 26 | google.joinpath("othermod.py").write_text("") 27 | env = dict(os.environ, PYTHONPATH=str(tmp_path)) 28 | cmd = [sys.executable, "-m", "google.othermod"] 29 | subprocess.check_call(cmd, env=env) 30 | 31 | # The ``google.cloud`` namespace package should not be masked by the presence of 32 | # ``google-cloud-alloydb`` and ``google-cloud-alloydb-connector``. 33 | google_cloud = tmp_path / "google" / "cloud" 34 | google_cloud.mkdir() 35 | google_cloud.joinpath("othermod.py").write_text("") 36 | env = dict(os.environ, PYTHONPATH=str(tmp_path)) 37 | cmd = [sys.executable, "-m", "google.cloud.othermod"] 38 | subprocess.check_call(cmd, env=env) 39 | -------------------------------------------------------------------------------- /tests/unit/test_rate_limiter.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import asyncio 16 | 17 | import pytest 18 | 19 | from google.cloud.alloydbconnector.rate_limiter import AsyncRateLimiter 20 | 21 | 22 | @pytest.mark.asyncio 23 | async def test_rate_limiter_throttles_requests() -> None: 24 | """Test to check whether rate limiter will throttle incoming requests.""" 25 | counter = 0 26 | # allow 2 requests to go through every 10 seconds 27 | rate_limiter = AsyncRateLimiter(max_capacity=2, rate=1 / 10) 28 | 29 | async def increment() -> None: 30 | await rate_limiter.acquire() 31 | nonlocal counter 32 | counter += 1 33 | 34 | # create 5 tasks calling increment() 35 | tasks = [asyncio.create_task(increment()) for _ in range(5)] 36 | 37 | # wait 5 seconds and check tasks 38 | done, pending = await asyncio.wait(tasks, timeout=5) 39 | 40 | # verify 2 tasks completed and 3 pending due to rate limiter 41 | assert counter == 2 42 | assert len(done) == 2 43 | assert len(pending) == 3 44 | 45 | # cleanup pending tasks 46 | for task in pending: 47 | task.cancel() 48 | 49 | 50 | @pytest.mark.asyncio 51 | async def test_rate_limiter_completes_all_tasks() -> None: 52 | """Test to check all requests will go through rate limiter successfully.""" 53 | counter = 0 54 | # allow 1 request to go through per second 55 | rate_limiter = AsyncRateLimiter(max_capacity=1, rate=1) 56 | 57 | async def increment() -> None: 58 | await rate_limiter.acquire() 59 | nonlocal counter 60 | counter += 1 61 | 62 | # create 5 tasks calling increment() 63 | tasks = [asyncio.create_task(increment()) for _ in range(5)] 64 | 65 | done, pending = await asyncio.wait(tasks, timeout=6) 66 | 67 | # verify all tasks done and none pending 68 | assert counter == 5 69 | assert len(done) == 5 70 | assert len(pending) == 0 71 | -------------------------------------------------------------------------------- /tests/unit/test_refresh_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from datetime import datetime 16 | from datetime import timedelta 17 | from datetime import timezone 18 | 19 | import pytest 20 | 21 | from google.cloud.alloydbconnector.refresh_utils import _seconds_until_refresh 22 | 23 | 24 | def test_seconds_until_refresh_over_1_hour() -> None: 25 | """ 26 | Test _seconds_until_refresh returns proper time in seconds. 27 | If expiration is over 1 hour, should return duration/2. 28 | """ 29 | # using pytest.approx since sometimes can be off by a second 30 | assert ( 31 | pytest.approx( 32 | _seconds_until_refresh(datetime.now(timezone.utc) + timedelta(minutes=62)), 33 | 1, 34 | ) 35 | == 31 * 60 36 | ) 37 | 38 | 39 | def test_seconds_until_refresh_under_1_hour_over_4_mins() -> None: 40 | """ 41 | Test _seconds_until_refresh returns proper time in seconds. 42 | If expiration is under 1 hour and over 4 minutes, 43 | should return duration-refresh_buffer (refresh_buffer = 4 minutes). 44 | """ 45 | # using pytest.approx since sometimes can be off by a second 46 | assert ( 47 | pytest.approx( 48 | _seconds_until_refresh(datetime.now(timezone.utc) + timedelta(minutes=5)), 49 | 1, 50 | ) 51 | == 60 52 | ) 53 | 54 | 55 | def test_seconds_until_refresh_under_4_mins() -> None: 56 | """ 57 | Test _seconds_until_refresh returns proper time in seconds. 58 | If expiration is under 4 minutes, should return 0. 59 | """ 60 | assert ( 61 | _seconds_until_refresh(datetime.now(timezone.utc) + timedelta(minutes=3)) == 0 62 | ) 63 | -------------------------------------------------------------------------------- /tests/unit/test_static.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from mocks import FakeInstance 16 | from mocks import write_static_info 17 | 18 | from google.cloud.alloydbconnector.connection_info import ConnectionInfo 19 | from google.cloud.alloydbconnector.static import StaticConnectionInfoCache 20 | 21 | 22 | def test_StaticConnectionInfoCache_init() -> None: 23 | """ 24 | Test that StaticConnectionInfoCache.__init__ populates its ConnectionInfo 25 | object. 26 | """ 27 | i = FakeInstance() 28 | static_info = write_static_info(i) 29 | cache = StaticConnectionInfoCache(i.uri(), static_info) 30 | assert len(cache._info.cert_chain) == 3 31 | assert cache._info.ca_cert 32 | assert cache._info.key 33 | assert cache._info.ip_addrs == { 34 | "PRIVATE": i.ip_addrs["PRIVATE"], 35 | "PUBLIC": i.ip_addrs["PUBLIC"], 36 | "PSC": i.ip_addrs["PSC"], 37 | } 38 | assert cache._info.expiration 39 | 40 | 41 | def test_StaticConnectionInfoCache_init_trailing_dot_dns() -> None: 42 | """ 43 | Test that StaticConnectionInfoCache.__init__ populates its ConnectionInfo 44 | object correctly when its PSC DNS name contains a trailing dot. 45 | """ 46 | i = FakeInstance() 47 | no_trailing_dot_dns = i.ip_addrs["PSC"] 48 | i.ip_addrs["PSC"] += "." 49 | static_info = write_static_info(i) 50 | cache = StaticConnectionInfoCache(i.uri(), static_info) 51 | assert len(cache._info.cert_chain) == 3 52 | assert cache._info.ca_cert 53 | assert cache._info.key 54 | assert cache._info.ip_addrs == { 55 | "PRIVATE": i.ip_addrs["PRIVATE"], 56 | "PUBLIC": i.ip_addrs["PUBLIC"], 57 | "PSC": no_trailing_dot_dns, 58 | } 59 | assert cache._info.expiration 60 | 61 | 62 | async def test_StaticConnectionInfoCache_force_refresh() -> None: 63 | """ 64 | Test that StaticConnectionInfoCache.force_refresh is a no-op. 65 | """ 66 | i = FakeInstance() 67 | static_info = write_static_info(i) 68 | cache = StaticConnectionInfoCache(i.uri(), static_info) 69 | conn_info = cache._info 70 | await cache.force_refresh() 71 | conn_info2 = cache._info 72 | assert conn_info2 == conn_info 73 | 74 | 75 | async def test_StaticConnectionInfoCache_connect_info() -> None: 76 | """ 77 | Test that StaticConnectionInfoCache.connect_info returns the ConnectionInfo 78 | object. 79 | """ 80 | i = FakeInstance() 81 | static_info = write_static_info(i) 82 | cache = StaticConnectionInfoCache(i.uri(), static_info) 83 | # check that cached connection info is now set 84 | assert isinstance(cache._info, ConnectionInfo) 85 | conn_info = cache._info 86 | # check that calling connect_info uses cached info 87 | conn_info2 = await cache.connect_info() 88 | assert conn_info2 == conn_info 89 | 90 | 91 | async def test_StaticConnectionInfoCache_close() -> None: 92 | """ 93 | Test that StaticConnectionInfoCache.close is a no-op. 94 | """ 95 | i = FakeInstance() 96 | static_info = write_static_info(i) 97 | cache = StaticConnectionInfoCache(i.uri(), static_info) 98 | conn_info = cache._info 99 | await cache.close() 100 | conn_info2 = cache._info 101 | assert conn_info2 == conn_info 102 | -------------------------------------------------------------------------------- /tests/unit/test_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from google.cloud.alloydbconnector.utils import strip_http_prefix 16 | 17 | 18 | def test_strip_http_prefix_with_empty_url() -> None: 19 | assert strip_http_prefix("") == "" 20 | 21 | 22 | def test_strip_http_prefix_with_url_having_http_prefix() -> None: 23 | assert strip_http_prefix("http://google.com") == "google.com" 24 | 25 | 26 | def test_strip_http_prefix_with_url_having_https_prefix() -> None: 27 | assert strip_http_prefix("https://google.com") == "google.com" 28 | --------------------------------------------------------------------------------