├── .github └── workflows │ ├── publish.yml │ └── release.yml ├── .gitignore ├── CHANGELOG.md ├── LICENSE ├── Makefile ├── README.md ├── docker-compose.yml ├── poetry.lock ├── pyproject.toml ├── src └── harlequin_postgres │ ├── __init__.py │ ├── adapter.py │ ├── catalog.py │ ├── cli_options.py │ ├── completions.py │ ├── interactions.py │ ├── keywords.tsv │ ├── loaders.py │ └── py.typed └── tests ├── conftest.py ├── test_adapter.py └── test_catalog.py /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Build and Publish Package 2 | 3 | on: 4 | pull_request: 5 | branches: 6 | - main 7 | types: 8 | - closed 9 | 10 | jobs: 11 | publish-package: 12 | if: ${{ github.event.pull_request.merged == true && startsWith(github.event.pull_request.head.ref, 'release/v') }} 13 | runs-on: ubuntu-latest 14 | 15 | steps: 16 | - name: Check out harlequin-postgres main branch 17 | uses: actions/checkout@v4 18 | with: 19 | ref: main 20 | - name: Set up Python 3.10 21 | uses: actions/setup-python@v4 22 | with: 23 | python-version: "3.10" 24 | - name: Install Poetry 25 | uses: snok/install-poetry@v1 26 | with: 27 | version: 1.6.1 28 | - name: Configure poetry 29 | run: poetry config --no-interaction pypi-token.pypi ${{ secrets.HARLEQUIN_PG_PYPI_TOKEN }} 30 | - name: Get harlequin-postgres Version 31 | id: harlequin_pg_version 32 | run: echo "harlequin_pg_version=$(poetry version --short)" >> $GITHUB_OUTPUT 33 | - name: Build package 34 | run: poetry build --no-interaction 35 | - name: Publish package to PyPI 36 | run: poetry publish --no-interaction 37 | - name: Create a Github Release 38 | uses: softprops/action-gh-release@v1 39 | with: 40 | tag_name: v${{ steps.harlequin_pg_version.outputs.harlequin_pg_version }} 41 | target_commitish: main 42 | token: ${{ secrets.HARLEQUIN_PG_RELEASE_TOKEN }} 43 | body_path: CHANGELOG.md 44 | files: | 45 | LICENSE 46 | dist/*harlequin*.whl 47 | dist/*harlequin*.tar.gz 48 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Create Release Branch 2 | 3 | on: 4 | workflow_dispatch: 5 | inputs: 6 | newVersion: 7 | description: A version number for this release (e.g., "0.1.0") 8 | required: true 9 | 10 | jobs: 11 | prepare-release: 12 | runs-on: ubuntu-latest 13 | permissions: 14 | contents: write 15 | pull-requests: write 16 | 17 | steps: 18 | - name: Check out harlequin-postgres main branch 19 | uses: actions/checkout@v4 20 | with: 21 | ref: main 22 | - name: Set up Python 3.10 23 | uses: actions/setup-python@v4 24 | with: 25 | python-version: "3.10" 26 | - name: Install Poetry 27 | uses: snok/install-poetry@v1 28 | with: 29 | version: 1.6.1 30 | - name: Create release branch 31 | run: | 32 | git checkout -b release/v${{ github.event.inputs.newVersion }} 33 | git push --set-upstream origin release/v${{ github.event.inputs.newVersion }} 34 | - name: Bump version 35 | run: poetry version ${{ github.event.inputs.newVersion }} --no-interaction 36 | - name: Ensure package can be built 37 | run: poetry build --no-interaction 38 | - name: Update CHANGELOG 39 | uses: thomaseizinger/keep-a-changelog-new-release@v1 40 | with: 41 | version: ${{ github.event.inputs.newVersion }} 42 | - name: Commit Changes 43 | uses: stefanzweifel/git-auto-commit-action@v5 44 | with: 45 | commit_message: Bumps version to ${{ github.event.inputs.newVersion }} 46 | - name: Create pull request into main 47 | uses: thomaseizinger/create-pull-request@1.3.1 48 | env: 49 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 50 | with: 51 | head: release/v${{ github.event.inputs.newVersion }} 52 | base: main 53 | title: v${{ github.event.inputs.newVersion }} 54 | body: > 55 | This PR was automatically generated. It bumps the version number 56 | in pyproject.toml and updates CHANGELOG.md. You may have to close 57 | this PR and reopen it to get the required checks to run. 58 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | profile.html 2 | 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | *.py[cod] 6 | *$py.class 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Distribution / packaging 12 | Pipfile 13 | .Python 14 | build/ 15 | develop-eggs/ 16 | dist/ 17 | downloads/ 18 | eggs/ 19 | .eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | wheels/ 26 | share/python-wheels/ 27 | *.egg-info/ 28 | .installed.cfg 29 | *.egg 30 | MANIFEST 31 | 32 | # PyInstaller 33 | # Usually these files are written by a python script from a template 34 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 35 | *.manifest 36 | *.spec 37 | 38 | # Installer logs 39 | pip-log.txt 40 | pip-delete-this-directory.txt 41 | 42 | # Unit test / coverage reports 43 | htmlcov/ 44 | .tox/ 45 | .nox/ 46 | .coverage 47 | .coverage.* 48 | .cache 49 | nosetests.xml 50 | coverage.xml 51 | *.cover 52 | *.py,cover 53 | .hypothesis/ 54 | .pytest_cache/ 55 | cover/ 56 | 57 | # Translations 58 | *.mo 59 | *.pot 60 | 61 | # Django stuff: 62 | *.log 63 | local_settings.py 64 | db.sqlite3 65 | db.sqlite3-journal 66 | 67 | # Flask stuff: 68 | instance/ 69 | .webassets-cache 70 | 71 | # Scrapy stuff: 72 | .scrapy 73 | 74 | # Sphinx documentation 75 | docs/_build/ 76 | 77 | # PyBuilder 78 | .pybuilder/ 79 | target/ 80 | 81 | # Jupyter Notebook 82 | .ipynb_checkpoints 83 | 84 | # IPython 85 | profile_default/ 86 | ipython_config.py 87 | 88 | # pyenv 89 | # For a library or package, you might want to ignore these files since the code is 90 | # intended to run in multiple environments; otherwise, check them in: 91 | # .python-version 92 | 93 | # pipenv 94 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 95 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 96 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 97 | # install all needed dependencies. 98 | #Pipfile.lock 99 | 100 | # poetry 101 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 102 | # This is especially recommended for binary packages to ensure reproducibility, and is more 103 | # commonly ignored for libraries. 104 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 105 | #poetry.lock 106 | 107 | # pdm 108 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 109 | #pdm.lock 110 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 111 | # in version control. 112 | # https://pdm.fming.dev/#use-with-ide 113 | .pdm.toml 114 | 115 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 116 | __pypackages__/ 117 | 118 | # Celery stuff 119 | celerybeat-schedule 120 | celerybeat.pid 121 | 122 | # SageMath parsed files 123 | *.sage.py 124 | 125 | # Environments 126 | .env 127 | .venv 128 | env/ 129 | venv/ 130 | ENV/ 131 | env.bak/ 132 | venv.bak/ 133 | 134 | # Spyder project settings 135 | .spyderproject 136 | .spyproject 137 | 138 | # Rope project settings 139 | .ropeproject 140 | 141 | # mkdocs documentation 142 | /site 143 | 144 | # mypy 145 | .mypy_cache/ 146 | .dmypy.json 147 | dmypy.json 148 | 149 | # Pyre type checker 150 | .pyre/ 151 | 152 | # pytype static type analyzer 153 | .pytype/ 154 | 155 | # Cython debug symbols 156 | cython_debug/ 157 | 158 | # PyCharm 159 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 160 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 161 | # and can be added to the global gitignore or merged into this file. For a more nuclear 162 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 163 | #.idea/ 164 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Harlequin-Postgres CHANGELOG 2 | 3 | All notable changes to this project will be documented in this file. 4 | 5 | ## [Unreleased] 6 | 7 | ## [1.2.0] - 2025-02-27 8 | 9 | - Adds interactions to list relations, indexes, constraints, and to describe relations (similar to psql's `\d+`) ([tconbeer/harlequin#586](https://github.com/tconbeer/harlequin/discussions/586) - thank you [@JPFrancoia](https://github.com/JPFrancoia)!). 10 | 11 | ## [1.1.1] - 2025-02-05 12 | 13 | - This adapter now supports `infinity` and `-infinity` dates and timestamps by loading their values as `date[time].max` or `date[time].min` ([tconbeer/harlequin#690](https://github.com/tconbeer/harlequin/issues/690)). 14 | 15 | ## [1.1.0] - 2025-01-27 16 | 17 | - This adapter now lazy-loads the catalog, which will dramatically improve the catalog performance for large databases with thousands of objects. 18 | - This adapter now implements interactions for catalog items, like dropping tables, setting the search path, etc. 19 | 20 | ## [1.0.0] - 2025-01-07 21 | 22 | - Drops support for Python 3.8 23 | - Adds support for Python 3.13 24 | - Adds support for Harlequin 2.X 25 | 26 | ## [0.4.0] - 2024-08-20 27 | 28 | - Upgrades client library to `psycopg3` (from `psycopg2`). 29 | - Adds an implementation of `connection_id` to improve catalog and history persistence. 30 | - Implements `cancel()` to interrupt in-flight queries. 31 | 32 | ## [0.3.0] - 2024-07-22 33 | 34 | - Adds an implementation of `close` to gracefully close the connection pool on Harlequin shut-down. 35 | - Adds support for Harlequin Transaction Modes and manual transactions. 36 | 37 | ## [0.2.2] - 2024-01-09 38 | 39 | - Sorts databases, schemas, and relations alphabetically; sorts columns ordinally. ([#10](https://github.com/tconbeer/harlequin-postgres/issues/10) - thank you [@frankbreetz](https://github.com/frankbreetz)!) 40 | 41 | ## [0.2.1] - 2023-12-14 42 | 43 | - Lowercases inserted values for keyword completions. 44 | 45 | ## [0.2.0] - 2023-12-14 46 | 47 | ### Features 48 | 49 | - Implements get_completions for keywords, functions, and settings. 50 | 51 | ## [0.1.3] - 2023-11-28 52 | 53 | ### Bug fixes 54 | 55 | - Implements connection pools instead of sharing a connection across threads. 56 | 57 | ## [0.1.2] - 2023-11-27 58 | 59 | ### Bug fixes 60 | 61 | - Fixes issues with package metadata. 62 | 63 | ## [0.1.1] - 2023-11-27 64 | 65 | ### Bug fixes 66 | 67 | - Fixes typo in release script. 68 | 69 | ## [0.1.0] - 2023-11-27 70 | 71 | ### Features 72 | 73 | - Adds a basic Postgres adapter with most common connection options. 74 | 75 | [Unreleased]: https://github.com/tconbeer/harlequin-postgres/compare/1.2.0...HEAD 76 | 77 | [1.2.0]: https://github.com/tconbeer/harlequin-postgres/compare/1.1.1...1.2.0 78 | 79 | [1.1.1]: https://github.com/tconbeer/harlequin-postgres/compare/1.1.0...1.1.1 80 | 81 | [1.1.0]: https://github.com/tconbeer/harlequin-postgres/compare/1.0.0...1.1.0 82 | 83 | [1.0.0]: https://github.com/tconbeer/harlequin-postgres/compare/0.4.0...1.0.0 84 | 85 | [0.4.0]: https://github.com/tconbeer/harlequin-postgres/compare/0.3.0...0.4.0 86 | 87 | [0.3.0]: https://github.com/tconbeer/harlequin-postgres/compare/0.2.2...0.3.0 88 | 89 | [0.2.2]: https://github.com/tconbeer/harlequin-postgres/compare/0.2.1...0.2.2 90 | 91 | [0.2.1]: https://github.com/tconbeer/harlequin-postgres/compare/0.2.0...0.2.1 92 | 93 | [0.2.0]: https://github.com/tconbeer/harlequin-postgres/compare/0.1.3...0.2.0 94 | 95 | [0.1.3]: https://github.com/tconbeer/harlequin-postgres/compare/0.1.2...0.1.3 96 | 97 | [0.1.2]: https://github.com/tconbeer/harlequin-postgres/compare/0.1.1...0.1.2 98 | 99 | [0.1.1]: https://github.com/tconbeer/harlequin-postgres/compare/0.1.0...0.1.1 100 | 101 | [0.1.0]: https://github.com/tconbeer/harlequin-postgres/compare/8611e628dc9d28b6a24817c761cd8a6da11a87ad...0.1.0 102 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Ted Conbeer 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: check 2 | check: 3 | ruff format . 4 | ruff check . --fix 5 | mypy 6 | pytest 7 | 8 | .PHONY: init 9 | init: 10 | docker-compose up -d 11 | 12 | .PHONY: clean 13 | clean: 14 | docker-compose down 15 | 16 | .PHONY: serve 17 | serve: 18 | harlequin -P None -a postgres "postgresql://postgres:for-testing@localhost:5432/postgres" 19 | 20 | .PHONY: psql 21 | psql: 22 | PGPASSWORD=for-testing psql -h localhost -p 5432 -U postgres -E 23 | 24 | profile.html: $(wildcard src/**/*.py) 25 | pyinstrument -r html -o profile.html --from-path harlequin -a postgres "postgresql://postgres:for-testing@localhost:5432/postgres" 26 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # harlequin-postgres 2 | 3 | This project provides the Harlequin adapter for Postgres. For more information, see [harlequin.sh](https://harlequin.sh/docs/postgres/index). 4 | 5 | 6 | ## Installation 7 | 8 | You must install the `harlequin-postgres` package into the same environment as `harlequin`. The best and easiest way to do this is to use `uv` to install Harlequin with the `postgres` extra: 9 | 10 | ```bash 11 | uv tool install 'harlequin[postgres]' 12 | ``` 13 | 14 | ## Using Harlequin with Postgres 15 | 16 | To connect to a Postgres database, run Harlequin with the `-a postgres` option and pass a [Posgres DSN](https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING) as an argument: 17 | 18 | ```bash 19 | harlequin -a postgres "postgres://my-user:my-pass@localhost:5432/my-database" 20 | ``` 21 | 22 | ## Connection Options 23 | 24 | You can also pass all or parts of the connection string as separate options. The following is equivalent to the above DSN: 25 | 26 | ```bash 27 | harlequin -a postgres -h localhost -p 5432 -U my-user --password my-pass -d my-database 28 | ``` 29 | 30 | The supported connection options are: 31 | 32 | ``` 33 | host 34 | port 35 | dbname 36 | user 37 | password 38 | passfile 39 | require_auth 40 | channel_binding 41 | connect_timeout 42 | sslmode 43 | sslcert 44 | sslkey 45 | ``` 46 | 47 | For descriptions of each option, run: 48 | 49 | ``` 50 | harlequin --help 51 | ``` 52 | 53 | ## Environment Variables 54 | 55 | Harlequin's Postgres driver will load connection information from the standard `PG*` environment variables. Any options supplied at the command-line will override environment variables. 56 | 57 | 58 | ## Manual Transactions 59 | 60 | To use Manual transaction mode, click on the label in the Run Query Bar to toggle the transaction mode from Auto to Manual. 61 | 62 | ## Further Documentation 63 | 64 | For more information, see the [Harlequin Docs](https://harlequin.sh/docs/postgres/index). -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | services: 2 | 3 | db: 4 | image: postgres 5 | restart: always 6 | environment: 7 | POSTGRES_PASSWORD: for-testing 8 | volumes: 9 | - pgdata:/var/lib/postgresql/data 10 | ports: 11 | - 5432:5432 12 | 13 | volumes: 14 | pgdata: 15 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "harlequin-postgres" 3 | version = "1.2.0" 4 | description = "A Harlequin adapter for Postgres." 5 | authors = ["Ted Conbeer "] 6 | license = "MIT" 7 | readme = "README.md" 8 | packages = [ 9 | { include = "harlequin_postgres", from = "src" }, 10 | ] 11 | 12 | [tool.poetry.plugins."harlequin.adapter"] 13 | postgres = "harlequin_postgres:HarlequinPostgresAdapter" 14 | 15 | [tool.poetry.dependencies] 16 | python = ">=3.9,<3.14" 17 | harlequin = ">=1.25,<3" 18 | psycopg = { version = "^3.2", extras = ["binary", "pool"]} 19 | 20 | [tool.poetry.group.dev.dependencies] 21 | ruff = "^0.5" 22 | pytest = "^7.4.3" 23 | mypy = "^1.10.0" 24 | pre-commit = "^3.5.0" 25 | importlib_metadata = { version = ">=4.6.0", python = "<3.10.0" } 26 | pyinstrument = "^4.6.1" 27 | 28 | [build-system] 29 | requires = ["poetry-core"] 30 | build-backend = "poetry.core.masonry.api" 31 | 32 | 33 | [tool.ruff] 34 | target-version = "py39" 35 | 36 | [tool.ruff.lint] 37 | select = ["A", "B", "E", "F", "I"] 38 | 39 | [tool.mypy] 40 | python_version = "3.9" 41 | files = [ 42 | "src/**/*.py", 43 | "tests/**/*.py", 44 | ] 45 | mypy_path = "src:stubs" 46 | 47 | show_column_numbers = true 48 | 49 | # show error messages from unrelated files 50 | follow_imports = "normal" 51 | 52 | # be strict 53 | disallow_untyped_calls = true 54 | disallow_untyped_defs = true 55 | check_untyped_defs = true 56 | disallow_untyped_decorators = true 57 | disallow_incomplete_defs = true 58 | disallow_subclassing_any = true 59 | strict_optional = true 60 | 61 | warn_return_any = true 62 | warn_no_return = true 63 | warn_redundant_casts = true 64 | warn_unused_ignores = true 65 | warn_unused_configs = true 66 | 67 | no_implicit_reexport = true 68 | strict_equality = true 69 | -------------------------------------------------------------------------------- /src/harlequin_postgres/__init__.py: -------------------------------------------------------------------------------- 1 | from harlequin_postgres.adapter import HarlequinPostgresAdapter 2 | 3 | __all__ = ["HarlequinPostgresAdapter"] 4 | -------------------------------------------------------------------------------- /src/harlequin_postgres/adapter.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from itertools import cycle 4 | from typing import Any, Sequence 5 | 6 | from harlequin import ( 7 | HarlequinAdapter, 8 | HarlequinCompletion, 9 | HarlequinConnection, 10 | HarlequinCursor, 11 | HarlequinTransactionMode, 12 | ) 13 | from harlequin.catalog import Catalog, CatalogItem 14 | from harlequin.exception import HarlequinConnectionError, HarlequinQueryError 15 | from psycopg import Connection, Cursor, conninfo 16 | from psycopg.errors import QueryCanceled 17 | from psycopg.pq import TransactionStatus 18 | from psycopg_pool import ConnectionPool 19 | from textual_fastdatatable.backend import AutoBackendType 20 | 21 | from harlequin_postgres.catalog import DatabaseCatalogItem 22 | from harlequin_postgres.cli_options import POSTGRES_OPTIONS 23 | from harlequin_postgres.completions import _get_completions 24 | from harlequin_postgres.loaders import register_inf_loaders 25 | 26 | 27 | class HarlequinPostgresCursor(HarlequinCursor): 28 | def __init__(self, conn: HarlequinPostgresConnection, cur: Cursor) -> None: 29 | self.conn = conn 30 | self.cur = cur 31 | # we need to copy the description from the cursor in case the results are 32 | # fetched and the cursor is closed before columns() is called. 33 | assert cur.description is not None 34 | self.description = cur.description.copy() 35 | self._limit: int | None = None 36 | 37 | def columns(self) -> list[tuple[str, str]]: 38 | return [ 39 | (col.name, self.conn._short_column_type_from_oid(col.type_code)) 40 | for col in self.description 41 | ] 42 | 43 | def set_limit(self, limit: int) -> HarlequinPostgresCursor: 44 | self._limit = limit 45 | return self 46 | 47 | def fetchall(self) -> AutoBackendType: 48 | try: 49 | if self._limit is None: 50 | return self.cur.fetchall() 51 | else: 52 | return self.cur.fetchmany(self._limit) 53 | except QueryCanceled: 54 | return [] 55 | except Exception as e: 56 | raise HarlequinQueryError( 57 | msg=f"{e.__class__.__name__}: {e}", 58 | title="Harlequin encountered an error while executing your query.", 59 | ) from e 60 | finally: 61 | self.cur.close() 62 | 63 | 64 | class HarlequinPostgresConnection(HarlequinConnection): 65 | def __init__( 66 | self, 67 | conn_str: Sequence[str], 68 | *_: Any, 69 | init_message: str = "", 70 | options: dict[str, Any], 71 | ) -> None: 72 | self.init_message = init_message 73 | try: 74 | self.conn_info = conninfo.conninfo_to_dict( 75 | conninfo=conn_str[0] if conn_str else "", **options 76 | ) 77 | except Exception as e: 78 | raise HarlequinConnectionError( 79 | msg=str(e), 80 | title=( 81 | "Harlequin could not connect to Postgres. " 82 | "Invalid connection string." 83 | ), 84 | ) from e 85 | try: 86 | raw_timeout = self.conn_info.get("connect_timeout") 87 | timeout = float(raw_timeout) if raw_timeout is not None else 30.0 88 | except (TypeError, ValueError) as e: 89 | raise HarlequinConnectionError( 90 | msg=str(e), 91 | title=( 92 | "Harlequin could not connect to Postgres. " 93 | "Invalid value for connection_timeout." 94 | ), 95 | ) from e 96 | try: 97 | self.pool: ConnectionPool = ConnectionPool( 98 | conninfo=conn_str[0] if conn_str and conn_str[0] else "", 99 | min_size=2, 100 | max_size=5, 101 | kwargs=options, 102 | open=True, 103 | timeout=timeout, 104 | ) 105 | self._main_conn: Connection = self.pool.getconn() 106 | except Exception as e: 107 | raise HarlequinConnectionError( 108 | msg=str(e), title="Harlequin could not connect to Postgres." 109 | ) from e 110 | 111 | self._transaction_modes = cycle( 112 | [ 113 | HarlequinTransactionMode(label="Auto"), 114 | HarlequinTransactionMode( 115 | label="Manual", 116 | commit=self.commit, 117 | rollback=self.rollback, 118 | ), 119 | ] 120 | ) 121 | self.toggle_transaction_mode() 122 | 123 | def execute(self, query: str) -> HarlequinCursor | None: 124 | if ( 125 | self.transaction_mode.label != "Auto" 126 | and self._main_conn.info.transaction_status == TransactionStatus.IDLE 127 | ): 128 | cur = self._main_conn.cursor() 129 | cur.execute(query="begin;") 130 | cur.close() 131 | 132 | try: 133 | cur = self._main_conn.cursor() 134 | cur.execute(query=query) 135 | except QueryCanceled: 136 | cur.close() 137 | return None 138 | except Exception as e: 139 | cur.close() 140 | self.rollback() 141 | raise HarlequinQueryError( 142 | msg=str(e), 143 | title="Harlequin encountered an error while executing your query.", 144 | ) from e 145 | else: 146 | if cur.description is not None: 147 | return HarlequinPostgresCursor(self, cur) 148 | else: 149 | cur.close() 150 | return None 151 | 152 | def cancel(self) -> None: 153 | self._main_conn.cancel_safe() 154 | 155 | def commit(self) -> None: 156 | self._main_conn.commit() 157 | 158 | def rollback(self) -> None: 159 | self._main_conn.rollback() 160 | 161 | def get_catalog(self) -> Catalog: 162 | databases = self._get_databases() 163 | db_items: list[CatalogItem] = [ 164 | DatabaseCatalogItem.from_label(label=db, connection=self) 165 | for (db,) in databases 166 | ] 167 | return Catalog(items=db_items) 168 | 169 | def get_completions(self) -> list[HarlequinCompletion]: 170 | conn: Connection = self.pool.getconn() 171 | completions = _get_completions(conn) 172 | self.pool.putconn(conn) 173 | return completions 174 | 175 | def close(self) -> None: 176 | self.pool.putconn(self._main_conn) 177 | self.pool.close() 178 | 179 | @property 180 | def transaction_mode(self) -> HarlequinTransactionMode: 181 | return self._transaction_mode 182 | 183 | def toggle_transaction_mode(self) -> HarlequinTransactionMode: 184 | self._transaction_mode = next(self._transaction_modes) 185 | self._sync_transaction_mode() 186 | return self._transaction_mode 187 | 188 | def _sync_transaction_mode(self) -> None: 189 | """ 190 | Sync this class's transaction mode with the main connection 191 | """ 192 | conn = self._main_conn 193 | if self.transaction_mode.label == "Auto": 194 | conn.autocommit = True 195 | conn.commit() 196 | else: 197 | conn.autocommit = False 198 | 199 | def _get_databases(self) -> list[tuple[str]]: 200 | conn: Connection = self.pool.getconn() 201 | with conn.cursor() as cur: 202 | cur.execute( 203 | """ 204 | select datname 205 | from pg_database 206 | where 207 | datistemplate is false 208 | and datallowconn is true 209 | order by datname asc 210 | ;""" 211 | ) 212 | results: list[tuple[str]] = cur.fetchall() 213 | self.pool.putconn(conn) 214 | return results 215 | 216 | def _get_schemas(self, dbname: str) -> list[tuple[str]]: 217 | conn: Connection = self.pool.getconn() 218 | with conn.cursor() as cur: 219 | cur.execute( 220 | f""" 221 | select schema_name 222 | from information_schema.schemata 223 | where 224 | catalog_name = '{dbname}' 225 | and schema_name != 'information_schema' 226 | and schema_name not like 'pg_%' 227 | order by schema_name asc 228 | ;""" 229 | ) 230 | results: list[tuple[str]] = cur.fetchall() 231 | self.pool.putconn(conn) 232 | return results 233 | 234 | def _get_relations(self, dbname: str, schema: str) -> list[tuple[str, str]]: 235 | conn: Connection = self.pool.getconn() 236 | with conn.cursor() as cur: 237 | cur.execute( 238 | f""" 239 | select table_name, table_type 240 | from information_schema.tables 241 | where 242 | table_catalog = '{dbname}' 243 | and table_schema = '{schema}' 244 | order by table_name asc 245 | ;""" 246 | ) 247 | results: list[tuple[str, str]] = cur.fetchall() 248 | self.pool.putconn(conn) 249 | return results 250 | 251 | def _get_columns( 252 | self, dbname: str, schema: str, relation: str 253 | ) -> list[tuple[str, str]]: 254 | conn: Connection = self.pool.getconn() 255 | with conn.cursor() as cur: 256 | cur.execute( 257 | f""" 258 | select column_name, data_type 259 | from information_schema.columns 260 | where 261 | table_catalog = '{dbname}' 262 | and table_schema = '{schema}' 263 | and table_name = '{relation}' 264 | order by ordinal_position asc 265 | ;""" 266 | ) 267 | results: list[tuple[str, str]] = cur.fetchall() 268 | self.pool.putconn(conn) 269 | return results 270 | 271 | @staticmethod 272 | def _short_column_type(type_name: str) -> str: 273 | MAPPING = { 274 | "bigint": "##", 275 | "bigserial": "##", 276 | "bit": "010", 277 | "boolean": "t/f", 278 | "box": "□", 279 | "bytea": "b", 280 | "character": "s", 281 | "cidr": "ip", 282 | "circle": "○", 283 | "date": "d", 284 | "double": "#.#", 285 | "inet": "ip", 286 | "integer": "#", 287 | "interval": "|-|", 288 | "json": "{}", 289 | "jsonb": "b{}", 290 | "line": "—", 291 | "lseg": "-", 292 | "macaddr": "mac", 293 | "macaddr8": "mac", 294 | "money": "$$", 295 | "numeric": "#.#", 296 | "path": "╭", 297 | "pg_lsn": "lsn", 298 | "pg_snapshot": "snp", 299 | "point": "•", 300 | "polygon": "▽", 301 | "real": "#.#", 302 | "smallint": "#", 303 | "smallserial": "#", 304 | "serial": "#", 305 | "text": "s", 306 | "time": "t", 307 | "timestamp": "ts", 308 | "tsquery": "tsq", 309 | "tsvector": "tsv", 310 | "txid_snapshot": "snp", 311 | "uuid": "uid", 312 | "xml": "xml", 313 | "array": "[]", 314 | } 315 | return MAPPING.get(type_name.split("(")[0].split(" ")[0], "?") 316 | 317 | @staticmethod 318 | def _short_column_type_from_oid(oid: int) -> str: 319 | MAPPING = { 320 | 16: "t/f", 321 | 17: "b", 322 | 18: "s", 323 | 19: "s", 324 | 20: "##", 325 | 21: "#", 326 | 22: "[#]", 327 | 23: "#", 328 | 25: "s", 329 | 26: "oid", 330 | 114: "{}", 331 | 142: "xml", 332 | 600: "•", 333 | 601: "-", 334 | 602: "╭", 335 | 603: "□", 336 | 604: "▽", 337 | 628: "—", 338 | 651: "[ip]", 339 | 700: "#.#", 340 | 701: "#.#", 341 | 704: "|-|", 342 | 718: "○", 343 | 790: "$$", 344 | 829: "mac", 345 | 869: "ip", 346 | 650: "ip", 347 | 774: "mac", 348 | 1000: "[t/f]", 349 | 1001: "[b]", 350 | 1002: "[s]", 351 | 1003: "[s]", 352 | 1009: "[s]", 353 | 1013: "[oid]", 354 | 1014: "[s]", 355 | 1015: "[s]", 356 | 1016: "[#]", 357 | 1021: "[#.#]", 358 | 1022: "[#.#]", 359 | 1028: "[oid]", 360 | 1040: "[mac]", 361 | 1041: "[ip]", 362 | 1042: "s", 363 | 1043: "s", 364 | 1082: "d", 365 | 1083: "t", 366 | 1114: "ts", 367 | 1115: "[ts]", 368 | 1182: "[d]", 369 | 1183: "[t]", 370 | 1184: "ts", 371 | 1185: "[ts]", 372 | 1186: "|-|", 373 | 1187: "[|-|]", 374 | 1231: "[#.#]", 375 | 1266: "t", 376 | 1270: "[t]", 377 | 1560: "010", 378 | 1562: "010", 379 | 1700: "#.#", 380 | 2950: "uid", 381 | 3614: "tsv", 382 | 3615: "tsq", 383 | 3802: "b{}", 384 | } 385 | return MAPPING.get(oid, "?") 386 | 387 | 388 | class HarlequinPostgresAdapter(HarlequinAdapter): 389 | ADAPTER_OPTIONS = POSTGRES_OPTIONS 390 | IMPLEMENTS_CANCEL = True 391 | 392 | def __init__( 393 | self, 394 | conn_str: Sequence[str], 395 | host: str | None = None, 396 | port: str | None = None, 397 | dbname: str | None = None, 398 | user: str | None = None, 399 | password: str | None = None, 400 | passfile: str | None = None, 401 | require_auth: str | None = None, 402 | channel_binding: str | None = None, 403 | connect_timeout: int | float | None = None, 404 | sslmode: str | None = None, 405 | sslcert: str | None = None, 406 | sslkey: str | None = None, 407 | **_: Any, 408 | ) -> None: 409 | self.conn_str = conn_str 410 | self.options: dict[str, str | int | None] = { 411 | "host": host, 412 | "port": port, 413 | "dbname": dbname, 414 | "user": user, 415 | "password": password, 416 | "passfile": passfile, 417 | "require_auth": require_auth, 418 | "channel_binding": channel_binding, 419 | "connect_timeout": connect_timeout, # type: ignore[dict-item] 420 | "sslmode": sslmode, 421 | "sslcert": sslcert, 422 | "sslkey": sslkey, 423 | } 424 | 425 | @property 426 | def connection_id(self) -> str | None: 427 | """ 428 | Use a simplified connection string, with only the host, port, and database 429 | """ 430 | try: 431 | conn_info = conninfo.conninfo_to_dict( 432 | conninfo=self.conn_str[0] if self.conn_str else "", 433 | **self.options, 434 | ) 435 | except Exception: 436 | return None 437 | 438 | host = conn_info.get("host", "localhost") 439 | port = conn_info.get("port", "5432") 440 | dbname = conn_info.get("dbname", "postgres") 441 | return f"{host}:{port}/{dbname}" 442 | 443 | def connect(self) -> HarlequinPostgresConnection: 444 | if len(self.conn_str) > 1: 445 | raise HarlequinConnectionError( 446 | "Cannot provide multiple connection strings to the Postgres adapter. " 447 | f"{self.conn_str}" 448 | ) 449 | # before creating the connection, register updated type adapters, so 450 | # all subsequent connections will use those adapters 451 | register_inf_loaders() 452 | conn = HarlequinPostgresConnection(self.conn_str, options=self.options) 453 | return conn 454 | -------------------------------------------------------------------------------- /src/harlequin_postgres/catalog.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from dataclasses import dataclass 4 | from typing import TYPE_CHECKING 5 | 6 | from harlequin.catalog import InteractiveCatalogItem 7 | 8 | from harlequin_postgres.interactions import ( 9 | execute_drop_database_statement, 10 | execute_drop_foreign_table_statement, 11 | execute_drop_schema_statement, 12 | execute_drop_table_statement, 13 | execute_drop_view_statement, 14 | execute_use_statement, 15 | insert_columns_at_cursor, 16 | show_describe_relation, 17 | show_describe_table_constraints, 18 | show_describe_table_indexes, 19 | show_list_indexes, 20 | show_list_objects, 21 | show_select_star, 22 | show_view_definition, 23 | ) 24 | 25 | if TYPE_CHECKING: 26 | from harlequin_postgres.adapter import HarlequinPostgresConnection 27 | 28 | 29 | @dataclass 30 | class ColumnCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"]): 31 | parent: "RelationCatalogItem" | None = None 32 | 33 | @classmethod 34 | def from_parent( 35 | cls, 36 | parent: "RelationCatalogItem", 37 | label: str, 38 | type_label: str, 39 | ) -> "ColumnCatalogItem": 40 | column_qualified_identifier = f'{parent.qualified_identifier}."{label}"' 41 | column_query_name = f'"{label}"' 42 | return cls( 43 | qualified_identifier=column_qualified_identifier, 44 | query_name=column_query_name, 45 | label=label, 46 | type_label=type_label, 47 | connection=parent.connection, 48 | parent=parent, 49 | loaded=True, 50 | ) 51 | 52 | 53 | @dataclass 54 | class RelationCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"]): 55 | INTERACTIONS = [ 56 | ("Insert Columns at Cursor", insert_columns_at_cursor), 57 | ("Preview Data", show_select_star), 58 | ("Describe Relation (\\d+)", show_describe_relation), 59 | ] 60 | parent: "SchemaCatalogItem" | None = None 61 | 62 | def fetch_children(self) -> list[ColumnCatalogItem]: 63 | if self.parent is None or self.parent.parent is None or self.connection is None: 64 | return [] 65 | result = self.connection._get_columns( 66 | self.parent.parent.label, self.parent.label, self.label 67 | ) 68 | return [ 69 | ColumnCatalogItem.from_parent( 70 | parent=self, 71 | label=column_name, 72 | type_label=self.connection._short_column_type(column_type), 73 | ) 74 | for column_name, column_type in result 75 | ] 76 | 77 | 78 | class ViewCatalogItem(RelationCatalogItem): 79 | INTERACTIONS = RelationCatalogItem.INTERACTIONS + [ 80 | ("Show View Definition", show_view_definition), 81 | ("Drop View", execute_drop_view_statement), 82 | ] 83 | 84 | @classmethod 85 | def from_parent( 86 | cls, 87 | parent: "SchemaCatalogItem", 88 | label: str, 89 | ) -> "ViewCatalogItem": 90 | relation_query_name = f'"{parent.label}"."{label}"' 91 | relation_qualified_identifier = f'{parent.qualified_identifier}."{label}"' 92 | return cls( 93 | qualified_identifier=relation_qualified_identifier, 94 | query_name=relation_query_name, 95 | label=label, 96 | type_label="v", 97 | connection=parent.connection, 98 | parent=parent, 99 | ) 100 | 101 | 102 | class TableCatalogItem(RelationCatalogItem): 103 | INTERACTIONS = RelationCatalogItem.INTERACTIONS + [ 104 | ("Describe Indexes", show_describe_table_indexes), 105 | ("Describe Constraints", show_describe_table_constraints), 106 | ("Drop Table", execute_drop_table_statement), 107 | ] 108 | 109 | @classmethod 110 | def from_parent( 111 | cls, 112 | parent: "SchemaCatalogItem", 113 | label: str, 114 | ) -> "TableCatalogItem": 115 | relation_query_name = f'"{parent.label}"."{label}"' 116 | relation_qualified_identifier = f'{parent.qualified_identifier}."{label}"' 117 | return cls( 118 | qualified_identifier=relation_qualified_identifier, 119 | query_name=relation_query_name, 120 | label=label, 121 | type_label="t", 122 | connection=parent.connection, 123 | parent=parent, 124 | ) 125 | 126 | 127 | class TempTableCatalogItem(TableCatalogItem): 128 | @classmethod 129 | def from_parent( 130 | cls, 131 | parent: "SchemaCatalogItem", 132 | label: str, 133 | ) -> "TempTableCatalogItem": 134 | relation_query_name = f'"{parent.label}"."{label}"' 135 | relation_qualified_identifier = f'{parent.qualified_identifier}."{label}"' 136 | return cls( 137 | qualified_identifier=relation_qualified_identifier, 138 | query_name=relation_query_name, 139 | label=label, 140 | type_label="tmp", 141 | connection=parent.connection, 142 | parent=parent, 143 | ) 144 | 145 | 146 | class ForeignCatalogItem(TableCatalogItem): 147 | INTERACTIONS = RelationCatalogItem.INTERACTIONS + [ 148 | ("Drop Table", execute_drop_foreign_table_statement), 149 | ] 150 | 151 | @classmethod 152 | def from_parent( 153 | cls, 154 | parent: "SchemaCatalogItem", 155 | label: str, 156 | ) -> "ForeignCatalogItem": 157 | relation_query_name = f'"{parent.label}"."{label}"' 158 | relation_qualified_identifier = f'{parent.qualified_identifier}."{label}"' 159 | return cls( 160 | qualified_identifier=relation_qualified_identifier, 161 | query_name=relation_query_name, 162 | label=label, 163 | type_label="f", 164 | connection=parent.connection, 165 | parent=parent, 166 | ) 167 | 168 | 169 | @dataclass 170 | class SchemaCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"]): 171 | INTERACTIONS = [ 172 | ("Set Search Path", execute_use_statement), 173 | ("List Relations (\\d+)", show_list_objects), 174 | ("List Indexes (\\di+)", show_list_indexes), 175 | ("Drop Schema", execute_drop_schema_statement), 176 | ] 177 | parent: "DatabaseCatalogItem" | None = None 178 | 179 | @classmethod 180 | def from_parent( 181 | cls, 182 | parent: "DatabaseCatalogItem", 183 | label: str, 184 | ) -> "SchemaCatalogItem": 185 | schema_identifier = f'"{label}"' 186 | return cls( 187 | qualified_identifier=schema_identifier, 188 | query_name=schema_identifier, 189 | label=label, 190 | type_label="sch", 191 | connection=parent.connection, 192 | parent=parent, 193 | ) 194 | 195 | def fetch_children(self) -> list[RelationCatalogItem]: 196 | if self.parent is None or self.connection is None: 197 | return [] 198 | children: list[RelationCatalogItem] = [] 199 | result = self.connection._get_relations(self.parent.label, self.label) 200 | for table_label, table_type in result: 201 | if table_type == "VIEW": 202 | children.append( 203 | ViewCatalogItem.from_parent( 204 | parent=self, 205 | label=table_label, 206 | ) 207 | ) 208 | elif table_type == "LOCAL TEMPORARY": 209 | children.append( 210 | TempTableCatalogItem.from_parent( 211 | parent=self, 212 | label=table_label, 213 | ) 214 | ) 215 | elif table_type == "FOREIGN": 216 | children.append( 217 | ForeignCatalogItem.from_parent( 218 | parent=self, 219 | label=table_label, 220 | ) 221 | ) 222 | else: 223 | children.append( 224 | TableCatalogItem.from_parent( 225 | parent=self, 226 | label=table_label, 227 | ) 228 | ) 229 | 230 | return children 231 | 232 | 233 | class DatabaseCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"]): 234 | INTERACTIONS = [ 235 | ("List Relations (\\d+)", show_list_objects), 236 | ("List Indexes (\\di+)", show_list_indexes), 237 | ("Drop Database", execute_drop_database_statement), 238 | ] 239 | 240 | @classmethod 241 | def from_label( 242 | cls, label: str, connection: "HarlequinPostgresConnection" 243 | ) -> "DatabaseCatalogItem": 244 | database_identifier = f'"{label}"' 245 | return cls( 246 | qualified_identifier=database_identifier, 247 | query_name=database_identifier, 248 | label=label, 249 | type_label="db", 250 | connection=connection, 251 | ) 252 | 253 | def fetch_children(self) -> list[SchemaCatalogItem]: 254 | if self.connection is None: 255 | return [] 256 | schemas = self.connection._get_schemas(self.label) 257 | return [ 258 | SchemaCatalogItem.from_parent( 259 | parent=self, 260 | label=schema_label, 261 | ) 262 | for (schema_label,) in schemas 263 | ] 264 | -------------------------------------------------------------------------------- /src/harlequin_postgres/cli_options.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from harlequin.options import ( 4 | FlagOption, # noqa 5 | ListOption, # noqa 6 | PathOption, # noqa 7 | SelectOption, # noqa 8 | TextOption, 9 | ) 10 | 11 | host = TextOption( 12 | name="host", 13 | description=( 14 | "Specifies the host name of the machine on which the server is running. " 15 | "If the value begins with a slash, it is used as the directory for the " 16 | "Unix-domain socket." 17 | ), 18 | short_decls=["-h"], 19 | default="localhost", 20 | ) 21 | 22 | 23 | port = TextOption( 24 | name="port", 25 | description=( 26 | "Port number to connect to at the server host, or socket file name extension " 27 | "for Unix-domain connections." 28 | ), 29 | short_decls=["-p"], 30 | default="5432", 31 | ) 32 | 33 | 34 | dbname = TextOption( 35 | name="dbname", 36 | description=("The database name to use when connecting with the Postgres server."), 37 | short_decls=["-d"], 38 | default="postgres", 39 | ) 40 | 41 | 42 | user = TextOption( 43 | name="user", 44 | description=("PostgreSQL user name to connect as."), 45 | short_decls=["-u", "--username", "-U"], 46 | ) 47 | 48 | 49 | password = TextOption( 50 | name="password", 51 | description=("Password to be used if the server demands password authentication."), 52 | ) 53 | 54 | 55 | passfile = PathOption( 56 | name="passfile", 57 | description=( 58 | "Specifies the name of the file used to store passwords. Defaults to " 59 | r"~/.pgpass, or %APPDATA%\postgresql\pgpass.conf on Windows. (No error is " 60 | "reported if this file does not exist.)" 61 | ), 62 | resolve_path=True, 63 | exists=False, 64 | file_okay=True, 65 | dir_okay=False, 66 | ) 67 | 68 | require_auth = SelectOption( 69 | name="require_auth", 70 | description=( 71 | "Specifies the authentication method that the client requires from the server. " 72 | "If the server does not use the required method to authenticate the client, or " 73 | "if the authentication handshake is not fully completed by the server, the " 74 | "connection will fail." 75 | ), 76 | choices=["password", "md5", "gss", "sspi", "scram-sha-256", "none"], 77 | ) 78 | 79 | channel_binding = SelectOption( 80 | name="channel_binding", 81 | description=( 82 | "This option controls the client's use of channel binding. A setting of " 83 | "require means that the connection must employ channel binding, prefer " 84 | "means that the client will choose channel binding if available, and " 85 | "disable prevents the use of channel binding. The default is prefer if " 86 | "PostgreSQL is compiled with SSL support; otherwise the default is disable." 87 | ), 88 | choices=["require", "prefer", "disable"], 89 | ) 90 | 91 | 92 | def _int_validator(s: str | None) -> tuple[bool, str]: 93 | if s is None: 94 | return True, "" 95 | try: 96 | _ = int(s) 97 | except ValueError: 98 | return False, f"Cannot convert {s} to an int!" 99 | else: 100 | return True, "" 101 | 102 | 103 | connect_timeout = TextOption( 104 | name="connect_timeout", 105 | description=( 106 | "Maximum time to wait while connecting, in seconds (write as an integer, " 107 | "e.g., 10)." 108 | ), 109 | validator=_int_validator, 110 | ) 111 | 112 | sslmode = SelectOption( 113 | name="sslmode", 114 | description=( 115 | "Determines whether or with what priority a secure SSL TCP/IP connection will " 116 | "be negotiated with the server." 117 | ), 118 | choices=["disable", "allow", "prefer", "require", "verify-ca", "verify-full"], 119 | default="prefer", 120 | ) 121 | 122 | sslcert = PathOption( 123 | name="sslcert", 124 | description=( 125 | "Specifies the file name of the client SSL certificate. " 126 | "Ignored if an SSL connection is not made." 127 | ), 128 | default="~/.postgresql/postgresql.crt", 129 | ) 130 | 131 | sslkey = TextOption( 132 | name="sslkey", 133 | description=( 134 | "Specifies the location for the secret key used for the client certificate. " 135 | "It can either specify a file name that will be used instead of the default " 136 | "~/.postgresql/postgresql.key, or it can specify a key obtained from an " 137 | "external engine. An external engine specification should consist of a " 138 | "colon-separated engine name and an engine-specific key identifier. This " 139 | "parameter is ignored if an SSL connection is not made." 140 | ), 141 | ) 142 | 143 | 144 | POSTGRES_OPTIONS = [ 145 | host, 146 | port, 147 | dbname, 148 | user, 149 | password, 150 | passfile, 151 | require_auth, 152 | channel_binding, 153 | connect_timeout, 154 | sslmode, 155 | sslcert, 156 | sslkey, 157 | ] 158 | -------------------------------------------------------------------------------- /src/harlequin_postgres/completions.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | import csv 4 | from pathlib import Path 5 | 6 | from harlequin import HarlequinCompletion 7 | from psycopg import Connection 8 | 9 | 10 | def _get_completions(conn: Connection) -> list[HarlequinCompletion]: 11 | completions: list[HarlequinCompletion] = [] 12 | 13 | # source: https://www.postgresql.org/docs/current/sql-keywords-appendix.html 14 | keyword_path = Path(__file__).parent / "keywords.tsv" 15 | with keyword_path.open("r") as f: 16 | keyword_reader = csv.reader( 17 | f, 18 | delimiter="\t", 19 | ) 20 | _header = next(keyword_reader) 21 | for keyword, kind, _, _, _ in keyword_reader: 22 | completions.append( 23 | HarlequinCompletion( 24 | label=keyword.lower(), 25 | type_label="kw", 26 | value=keyword.lower(), 27 | priority=100 if kind.startswith("reserved") else 1000, 28 | context=None, 29 | ) 30 | ) 31 | 32 | with conn.cursor() as cur: 33 | cur.execute( 34 | r""" 35 | select distinct 36 | routine_name as label, 37 | case when routine_type is null then 'agg' else 'fn' end as type_label, 38 | case 39 | when routine_schema = 'pg_catalog' -- 40 | then null 41 | else routine_schema 42 | end as context 43 | from information_schema.routines 44 | where 45 | length(routine_name) < 37 46 | and routine_name not ilike '\_%' 47 | and routine_name not ilike 'pg\_%' 48 | and routine_name not ilike 'binary\_upgrade\_%' 49 | 50 | ;""" 51 | ) 52 | results = cur.fetchall() 53 | for label, type_label, context in results: 54 | completions.append( 55 | HarlequinCompletion( 56 | label=label, 57 | type_label=type_label, 58 | value=label, 59 | priority=1000, 60 | context=context, 61 | ) 62 | ) 63 | 64 | with conn.cursor() as cur: 65 | cur.execute("""select distinct name as label from pg_settings""") 66 | results = cur.fetchall() 67 | for (label,) in results: 68 | completions.append( 69 | HarlequinCompletion( 70 | label=label, type_label="set", value=label, priority=2000, context=None 71 | ) 72 | ) 73 | 74 | return sorted(completions) 75 | -------------------------------------------------------------------------------- /src/harlequin_postgres/interactions.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from textwrap import dedent 4 | from typing import TYPE_CHECKING, Literal, Sequence 5 | 6 | from harlequin.catalog import CatalogItem 7 | from harlequin.exception import HarlequinQueryError 8 | 9 | if TYPE_CHECKING: 10 | from harlequin.driver import HarlequinDriver 11 | 12 | from harlequin_postgres.catalog import ( 13 | ColumnCatalogItem, 14 | DatabaseCatalogItem, 15 | RelationCatalogItem, 16 | SchemaCatalogItem, 17 | ViewCatalogItem, 18 | ) 19 | 20 | 21 | def execute_use_statement( 22 | item: "SchemaCatalogItem", 23 | driver: "HarlequinDriver", 24 | ) -> None: 25 | if item.connection is None: 26 | return 27 | try: 28 | item.connection.execute(f"set search_path to {item.qualified_identifier}") 29 | except HarlequinQueryError: 30 | driver.notify("Could not switch context", severity="error") 31 | raise 32 | else: 33 | driver.notify(f"Editor context switched to {item.label}") 34 | 35 | 36 | def execute_drop_schema_statement( 37 | item: "SchemaCatalogItem", 38 | driver: "HarlequinDriver", 39 | ) -> None: 40 | def _drop_schema() -> None: 41 | if item.connection is None: 42 | return 43 | try: 44 | item.connection.execute(f"drop schema {item.qualified_identifier} cascade") 45 | except HarlequinQueryError: 46 | driver.notify(f"Could not drop schema {item.label}", severity="error") 47 | raise 48 | else: 49 | driver.notify(f"Dropped schema {item.label}") 50 | driver.refresh_catalog() 51 | 52 | if item.children or item.fetch_children(): 53 | driver.confirm_and_execute(callback=_drop_schema) 54 | else: 55 | _drop_schema() 56 | 57 | 58 | def execute_drop_database_statement( 59 | item: "DatabaseCatalogItem", 60 | driver: "HarlequinDriver", 61 | ) -> None: 62 | def _drop_database() -> None: 63 | if item.connection is None: 64 | return 65 | try: 66 | item.connection.execute(f"drop database {item.qualified_identifier}") 67 | except HarlequinQueryError: 68 | driver.notify(f"Could not drop database {item.label}", severity="error") 69 | raise 70 | else: 71 | driver.notify(f"Dropped database {item.label}") 72 | driver.refresh_catalog() 73 | 74 | if item.children or item.fetch_children(): 75 | driver.confirm_and_execute(callback=_drop_database) 76 | else: 77 | _drop_database() 78 | 79 | 80 | def execute_drop_relation_statement( 81 | item: "RelationCatalogItem", 82 | driver: "HarlequinDriver", 83 | relation_type: Literal["view", "table", "foreign table"], 84 | ) -> None: 85 | def _drop_relation() -> None: 86 | if item.connection is None: 87 | return 88 | try: 89 | item.connection.execute(f"drop {relation_type} {item.qualified_identifier}") 90 | except HarlequinQueryError: 91 | driver.notify( 92 | f"Could not drop {relation_type} {item.label}", severity="error" 93 | ) 94 | raise 95 | else: 96 | driver.notify(f"Dropped {relation_type} {item.label}") 97 | driver.refresh_catalog() 98 | 99 | driver.confirm_and_execute(callback=_drop_relation) 100 | 101 | 102 | def execute_drop_table_statement( 103 | item: "RelationCatalogItem", driver: "HarlequinDriver" 104 | ) -> None: 105 | execute_drop_relation_statement(item=item, driver=driver, relation_type="table") 106 | 107 | 108 | def execute_drop_foreign_table_statement( 109 | item: "RelationCatalogItem", driver: "HarlequinDriver" 110 | ) -> None: 111 | execute_drop_relation_statement( 112 | item=item, driver=driver, relation_type="foreign table" 113 | ) 114 | 115 | 116 | def execute_drop_view_statement( 117 | item: "RelationCatalogItem", driver: "HarlequinDriver" 118 | ) -> None: 119 | execute_drop_relation_statement(item=item, driver=driver, relation_type="view") 120 | 121 | 122 | def show_select_star( 123 | item: "RelationCatalogItem", 124 | driver: "HarlequinDriver", 125 | ) -> None: 126 | driver.insert_text_in_new_buffer( 127 | dedent( 128 | f""" 129 | select * 130 | from {item.qualified_identifier} 131 | limit 100 132 | """.strip("\n") 133 | ) 134 | ) 135 | 136 | 137 | def show_list_objects( 138 | item: "SchemaCatalogItem" | "DatabaseCatalogItem", 139 | driver: "HarlequinDriver", 140 | ) -> None: 141 | # sourced from psql with -E, then the following command: 142 | # \dtvmsE+ .* 143 | 144 | # can't use isinstance due to circular reference 145 | if type(item).__name__ == "SchemaCatalogItem": 146 | where_clause = f"and n.nspname = '{item.label}'" 147 | else: 148 | where_clause = ( 149 | "and n.nspname not in ('pg_catalog', 'pg_toast', 'information_schema')" 150 | ) 151 | driver.insert_text_in_new_buffer( 152 | dedent( 153 | f""" 154 | select 155 | n.nspname as "Schema", 156 | c.relname as "Name", 157 | case c.relkind 158 | when 'r' then 'table' 159 | when 'v' then 'view' 160 | when 'm' then 'materialized view' 161 | when 'S' 162 | then 'sequence' 163 | when 't' then 'TOAST table' 164 | when 'f' then 'foreign table' 165 | when 'p' then 'partitioned table' 166 | end as "Type", 167 | pg_catalog.pg_get_userbyid(c.relowner) as "Owner", 168 | case c.relpersistence 169 | when 'p' then 'permanent' 170 | when 't' then 'temporary' 171 | when 'u' then 'unlogged' 172 | end as "Persistence", 173 | am.amname as "Access method", 174 | pg_catalog.pg_size_pretty(pg_catalog.pg_table_size(c.oid)) as "Size", 175 | pg_catalog.obj_description(c.oid, 'pg_class') as "Description" 176 | from pg_catalog.pg_class c 177 | left join pg_catalog.pg_namespace n on n.oid = c.relnamespace 178 | left join pg_catalog.pg_am am on am.oid = c.relam 179 | where 180 | c.relkind IN ('r','p','t','v','m', 's', 'S', 'f') 181 | {where_clause} 182 | order by 1,2; 183 | """.strip("\n") 184 | ) 185 | ) 186 | 187 | 188 | def show_list_indexes( 189 | item: "SchemaCatalogItem" | "DatabaseCatalogItem", 190 | driver: "HarlequinDriver", 191 | ) -> None: 192 | # sourced from psql with -E, then the following command: 193 | # \dis+ .* 194 | 195 | # can't use isinstance due to circular reference 196 | if type(item).__name__ == "SchemaCatalogItem": 197 | where_clause = f"and n.nspname = '{item.label}'" 198 | else: 199 | where_clause = ( 200 | "and n.nspname not in ('pg_catalog', 'pg_toast', 'information_schema')" 201 | ) 202 | driver.insert_text_in_new_buffer( 203 | dedent( 204 | f""" 205 | select 206 | n.nspname as "Schema", 207 | c.relname as "Name", 208 | case 209 | c.relkind 210 | when 'i' 211 | then 'index' 212 | when 'I' 213 | then 'partitioned index' 214 | end as "Type", 215 | pg_catalog.pg_get_userbyid(c.relowner) as "Owner", 216 | c2.relname as "Table", 217 | case 218 | c.relpersistence 219 | when 'p' 220 | then 'permanent' 221 | when 't' 222 | then 'temporary' 223 | when 'u' 224 | then 'unlogged' 225 | end as "Persistence", 226 | am.amname as "Access method", 227 | pg_catalog.pg_size_pretty(pg_catalog.pg_table_size(c.oid)) as "Size", 228 | pg_catalog.obj_description(c.oid, 'pg_class') as "Description" 229 | from pg_catalog.pg_class c 230 | left join pg_catalog.pg_namespace n on n.oid = c.relnamespace 231 | left join pg_catalog.pg_am am on am.oid = c.relam 232 | left join pg_catalog.pg_index i on i.indexrelid = c.oid 233 | left join pg_catalog.pg_class c2 on i.indrelid = c2.oid 234 | where 235 | c.relkind in ('i', 'I') 236 | {where_clause} 237 | order by 1,2; 238 | """.strip("\n") 239 | ) 240 | ) 241 | 242 | 243 | def show_describe_relation( 244 | item: "RelationCatalogItem", 245 | driver: "HarlequinDriver", 246 | ) -> None: 247 | # sourced from psql -E \d+ {my rel} 248 | # see https://stackoverflow.com/questions/60155968/using-results-of-d-command-in-psql 249 | if item.parent is None: 250 | driver.notify( 251 | f"Could not describe {item.label} due to missing schema reference.", 252 | severity="error", 253 | ) 254 | return 255 | driver.insert_text_in_new_buffer( 256 | dedent( 257 | f""" 258 | with 259 | index_columns as ( 260 | select i.indexrelid, c.oid as rel_oid, unnest(i.indkey) as attnum 261 | from pg_catalog.pg_index i 262 | join pg_catalog.pg_class c on c.oid = i.indrelid 263 | where c.relname = '{item.label}' 264 | ), 265 | index_column_counts as ( 266 | select rel_oid, attnum, count(*) as cnt 267 | from index_columns 268 | group by 1, 2 269 | ), 270 | constraint_columns as ( 271 | select con.oid, c.oid as rel_oid, unnest(con.conkey) as attnum 272 | from pg_catalog.pg_constraint con 273 | join pg_catalog.pg_class c on con.conrelid = c.oid 274 | where c.relname = '{item.label}' 275 | ), 276 | constraint_column_counts as ( 277 | select rel_oid, attnum, count(*) as cnt 278 | from constraint_columns 279 | group by 1, 2 280 | ), 281 | fkey_columns as ( 282 | select 283 | src.relname as src_name, 284 | src.relnamespace::regnamespace as src_schema, 285 | c.oid as rel_oid, 286 | unnest(con.confkey) as attnum 287 | from pg_catalog.pg_constraint con 288 | join pg_catalog.pg_class c on con.confrelid = c.oid 289 | join pg_catalog.pg_class src on con.conrelid = src.oid 290 | where c.relname = '{item.label}' 291 | ), 292 | fkey_references as ( 293 | select 294 | rel_oid, 295 | attnum, 296 | string_agg(src_schema || '.' || src_name, ', ') as sources 297 | from fkey_columns 298 | group by 1, 2 299 | ) 300 | select 301 | a.attname as "Column", 302 | pg_catalog.format_type(a.atttypid, a.atttypmod) as "Type", 303 | coll.collname as "Collation", 304 | case 305 | when a.attnotnull is true then 'not null' 306 | else '' 307 | end as "Nullable", 308 | pg_catalog.pg_get_expr(d.adbin, d.adrelid, true) as "Default", 309 | case 310 | a.attstorage 311 | when 'p' 312 | then 'plain' 313 | when 'x' 314 | then 'extended' 315 | when 'e' 316 | then 'external' 317 | when 'm' 318 | then 'main' 319 | else a.attstorage::text 320 | end as "Storage", 321 | case 322 | a.attcompression 323 | when 'p' 324 | then 'pglz' 325 | when 'l' 326 | then 'LZ4' 327 | else a.attcompression::text 328 | end as "Compression", 329 | case 330 | when a.attstattarget = -1 331 | then null 332 | else a.attstattarget 333 | end as "Stats target", 334 | case 335 | when index_column_counts.cnt > 0 then true else false 336 | end as "Has Index", 337 | case 338 | when constraint_column_counts.cnt > 0 then true else false 339 | end as "Has Constraint", 340 | fkey_references.sources as "Referenced by", 341 | pg_catalog.col_description(a.attrelid, a.attnum) as "Description" 342 | from pg_catalog.pg_attribute a 343 | join pg_catalog.pg_class c on a.attrelid = c.oid 344 | left join pg_catalog.pg_namespace n on n.oid = c.relnamespace 345 | left join pg_catalog.pg_collation coll on coll.oid = a.attcollation 346 | left join 347 | pg_catalog.pg_type t 348 | on (t.oid = a.atttypid and t.typcollation <> a.attcollation) 349 | left join 350 | pg_catalog.pg_attrdef d 351 | on (a.attrelid = d.adrelid and a.attnum = d.adnum and a.atthasdef) 352 | left join 353 | index_column_counts 354 | on a.attnum = index_column_counts.attnum 355 | and a.attrelid = index_column_counts.rel_oid 356 | left join 357 | constraint_column_counts 358 | on a.attnum = constraint_column_counts.attnum 359 | and a.attrelid = constraint_column_counts.rel_oid 360 | left join 361 | fkey_references 362 | on a.attnum = fkey_references.attnum 363 | and a.attrelid = fkey_references.rel_oid 364 | where 365 | c.relname = '{item.label}' 366 | and n.nspname = '{item.parent.label}' 367 | and a.attnum > 0 368 | and not a.attisdropped 369 | order by a.attnum 370 | """.strip("\n") 371 | ) 372 | ) 373 | 374 | 375 | def show_describe_table_indexes( 376 | item: "RelationCatalogItem", 377 | driver: "HarlequinDriver", 378 | ) -> None: 379 | if item.parent is None: 380 | driver.notify( 381 | f"Could not describe {item.label} due to missing schema reference.", 382 | severity="error", 383 | ) 384 | return 385 | driver.insert_text_in_new_buffer( 386 | dedent( 387 | f""" 388 | with 389 | index_columns as ( 390 | select 391 | i.indexrelid, c.oid as rel_oid, unnest(i.indkey) as attnum 392 | from pg_catalog.pg_index i 393 | join pg_catalog.pg_class c on c.oid = i.indrelid 394 | where c.relname = '{item.label}' 395 | ), 396 | index_column_names as ( 397 | select 398 | index_columns.indexrelid, 399 | string_agg(pg_attribute.attname, ', ') as columns 400 | from index_columns 401 | join 402 | pg_catalog.pg_attribute 403 | on index_columns.rel_oid = pg_attribute.attrelid 404 | and index_columns.attnum = pg_attribute.attnum 405 | group by 1 406 | ) 407 | 408 | select 409 | c.relname as "Table", 410 | i.indexrelid::regclass::text as "Index Name", 411 | pg_am.amname as "Index Type", 412 | index_column_names.columns as "Columns", 413 | i.indisprimary as "Is PK", 414 | i.indisunique as "Is Unique", 415 | i.indisclustered as "Is Clustered", 416 | i.indisvalid as "Is Valid" 417 | from pg_catalog.pg_index i 418 | join pg_catalog.pg_class c on c.oid = i.indrelid 419 | left join pg_catalog.pg_namespace n on n.oid = c.relnamespace 420 | join pg_catalog.pg_class ic on i.indexrelid = ic.oid 421 | join pg_catalog.pg_am on ic.relam = pg_am.oid 422 | left join 423 | index_column_names on i.indexrelid = index_column_names.indexrelid 424 | where 425 | c.relname = '{item.label}' 426 | and n.nspname = '{item.parent.label}' 427 | """.strip("\n") 428 | ) 429 | ) 430 | 431 | 432 | def show_describe_table_constraints( 433 | item: "RelationCatalogItem", 434 | driver: "HarlequinDriver", 435 | ) -> None: 436 | if item.parent is None: 437 | driver.notify( 438 | f"Could not describe {item.label} due to missing schema reference.", 439 | severity="error", 440 | ) 441 | return 442 | driver.insert_text_in_new_buffer( 443 | dedent( 444 | f""" 445 | with 446 | constraint_columns as ( 447 | select con.oid, c.oid as rel_oid, unnest(con.conkey) as attnum 448 | from pg_catalog.pg_constraint con 449 | join pg_catalog.pg_class c on con.conrelid = c.oid 450 | where c.relname = '{item.label}' 451 | ), 452 | constraint_column_names as ( 453 | select 454 | constraint_columns.oid, 455 | string_agg(pg_attribute.attname, ', ') as columns 456 | from constraint_columns 457 | join 458 | pg_catalog.pg_attribute 459 | on constraint_columns.rel_oid = pg_attribute.attrelid 460 | and constraint_columns.attnum = pg_attribute.attnum 461 | group by 1 462 | ), 463 | constraint_foreign_columns as ( 464 | select con.oid, c.oid as rel_oid, unnest(con.confkey) as attnum 465 | from pg_catalog.pg_constraint con 466 | join pg_catalog.pg_class c on con.conrelid = c.oid 467 | where c.relname = '{item.label}' 468 | ), 469 | constraint_foreign_column_names as ( 470 | select 471 | constraint_foreign_columns.oid, 472 | string_agg(pg_attribute.attname, ', ') as columns 473 | from constraint_foreign_columns 474 | join 475 | pg_catalog.pg_attribute 476 | on constraint_foreign_columns.rel_oid = pg_attribute.attrelid 477 | and constraint_foreign_columns.attnum = pg_attribute.attnum 478 | group by 1 479 | ) 480 | 481 | select 482 | c.relname as "Table", 483 | con.conname as "Constraint name", 484 | case 485 | con.contype 486 | when 'c' 487 | then 'Check' 488 | when 'n' 489 | then 'Not Null' 490 | when 'p' 491 | then 'Primary Key' 492 | when 'f' 493 | then 'Foreign Key' 494 | when 'u' 495 | then 'Unique' 496 | when 't' 497 | then 'Trigger' 498 | when 'x' 499 | then 'Exclusion' 500 | else con.contype::text 501 | end as "Constraint Type", 502 | constraint_column_names.columns as "Columns", 503 | c.relnamespace::regnamespace 504 | || '.' 505 | || fc.relname 506 | || '(' 507 | || constraint_foreign_column_names.columns 508 | || ')' as "References", 509 | con.conislocal as "Is Local", 510 | con.convalidated as "Is Validated", 511 | case 512 | con.confupdtype 513 | when 'a' 514 | then 'no action' 515 | when 'r' 516 | then 'restrict' 517 | when 'c' 518 | then 'cascade' 519 | when 'n' 520 | then 'set null' 521 | when 'd' 522 | then 'set default' 523 | else con.confupdtype::text 524 | end as "FK Update Type", 525 | case 526 | con.confdeltype 527 | when 'a' 528 | then 'no action' 529 | when 'r' 530 | then 'restrict' 531 | when 'c' 532 | then 'cascade' 533 | when 'n' 534 | then 'set null' 535 | when 'd' 536 | then 'set default' 537 | else con.confdeltype::text 538 | end as "FK Delete Type", 539 | case 540 | con.confmatchtype 541 | when 's' 542 | then 'simple' 543 | when 'f' 544 | then 'full' 545 | when 'p' 546 | then 'partial' 547 | else con.confmatchtype::text 548 | end as "FK Match Type" 549 | from pg_catalog.pg_constraint con 550 | join pg_catalog.pg_class c on con.conrelid = c.oid 551 | join pg_catalog.pg_namespace n on n.oid = c.relnamespace 552 | left join pg_catalog.pg_class fc on con.confrelid = fc.oid 553 | left join 554 | constraint_column_names on con.oid = constraint_column_names.oid 555 | left join 556 | constraint_foreign_column_names 557 | on con.oid = constraint_foreign_column_names.oid 558 | where 559 | c.relname = '{item.label}' 560 | and n.nspname = '{item.parent.label}' 561 | """.strip("\n") 562 | ) 563 | ) 564 | 565 | 566 | def show_view_definition( 567 | item: "ViewCatalogItem", 568 | driver: "HarlequinDriver", 569 | ) -> None: 570 | if item.connection is None or item.parent is None: 571 | return 572 | view_def_query = f""" 573 | select pg_catalog.pg_get_viewdef(c.oid, true) 574 | from pg_catalog.pg_class as c 575 | left join pg_catalog.pg_namespace n on n.oid = c.relnamespace 576 | where c.relname = '{item.label}' and n.nspname = '{item.parent.label}' 577 | """.strip("\n") 578 | cur = item.connection.execute(view_def_query) 579 | if cur is None: 580 | return 581 | result = cur.fetchall() 582 | if result is None: 583 | return 584 | view_def: str = result[0][0] 585 | driver.insert_text_in_new_buffer( 586 | f"-- View definition for {item.query_name}\n" + view_def 587 | ) 588 | 589 | 590 | def insert_columns_at_cursor( 591 | item: "RelationCatalogItem", 592 | driver: "HarlequinDriver", 593 | ) -> None: 594 | if item.loaded: 595 | cols: Sequence["CatalogItem" | "ColumnCatalogItem"] = item.children 596 | else: 597 | cols = item.fetch_children() 598 | driver.insert_text_at_selection(text=",\n".join(c.query_name for c in cols)) 599 | -------------------------------------------------------------------------------- /src/harlequin_postgres/keywords.tsv: -------------------------------------------------------------------------------- 1 | Key Word PostgreSQL SQL:2023 SQL:2016 SQL-92 2 | A non-reserved non-reserved 3 | ABORT non-reserved 4 | ABS reserved reserved 5 | ABSENT non-reserved reserved reserved 6 | ABSOLUTE non-reserved non-reserved non-reserved reserved 7 | ACCESS non-reserved 8 | ACCORDING non-reserved non-reserved 9 | ACOS reserved reserved 10 | ACTION non-reserved non-reserved non-reserved reserved 11 | ADA non-reserved non-reserved non-reserved 12 | ADD non-reserved non-reserved non-reserved reserved 13 | ADMIN non-reserved non-reserved non-reserved 14 | AFTER non-reserved non-reserved non-reserved 15 | AGGREGATE non-reserved 16 | ALL reserved reserved reserved reserved 17 | ALLOCATE reserved reserved reserved 18 | ALSO non-reserved 19 | ALTER non-reserved reserved reserved reserved 20 | ALWAYS non-reserved non-reserved non-reserved 21 | ANALYSE reserved 22 | ANALYZE reserved 23 | AND reserved reserved reserved reserved 24 | ANY reserved reserved reserved reserved 25 | ANY_VALUE reserved 26 | ARE reserved reserved reserved 27 | ARRAY reserved, requires AS reserved reserved 28 | ARRAY_AGG reserved reserved 29 | ARRAY_​MAX_​CARDINALITY reserved reserved 30 | AS reserved, requires AS reserved reserved reserved 31 | ASC reserved non-reserved non-reserved reserved 32 | ASENSITIVE non-reserved reserved reserved 33 | ASIN reserved reserved 34 | ASSERTION non-reserved non-reserved non-reserved reserved 35 | ASSIGNMENT non-reserved non-reserved non-reserved 36 | ASYMMETRIC reserved reserved reserved 37 | AT non-reserved reserved reserved reserved 38 | ATAN reserved reserved 39 | ATOMIC non-reserved reserved reserved 40 | ATTACH non-reserved 41 | ATTRIBUTE non-reserved non-reserved non-reserved 42 | ATTRIBUTES non-reserved non-reserved 43 | AUTHORIZATION reserved (can be function or type) reserved reserved reserved 44 | AVG reserved reserved reserved 45 | BACKWARD non-reserved 46 | BASE64 non-reserved non-reserved 47 | BEFORE non-reserved non-reserved non-reserved 48 | BEGIN non-reserved reserved reserved reserved 49 | BEGIN_FRAME reserved reserved 50 | BEGIN_PARTITION reserved reserved 51 | BERNOULLI non-reserved non-reserved 52 | BETWEEN non-reserved (cannot be function or type) reserved reserved reserved 53 | BIGINT non-reserved (cannot be function or type) reserved reserved 54 | BINARY reserved (can be function or type) reserved reserved 55 | BIT non-reserved (cannot be function or type) reserved 56 | BIT_LENGTH reserved 57 | BLOB reserved reserved 58 | BLOCKED non-reserved non-reserved 59 | BOM non-reserved non-reserved 60 | BOOLEAN non-reserved (cannot be function or type) reserved reserved 61 | BOTH reserved reserved reserved reserved 62 | BREADTH non-reserved non-reserved non-reserved 63 | BTRIM reserved 64 | BY non-reserved reserved reserved reserved 65 | C non-reserved non-reserved non-reserved 66 | CACHE non-reserved 67 | CALL non-reserved reserved reserved 68 | CALLED non-reserved reserved reserved 69 | CARDINALITY reserved reserved 70 | CASCADE non-reserved non-reserved non-reserved reserved 71 | CASCADED non-reserved reserved reserved reserved 72 | CASE reserved reserved reserved reserved 73 | CAST reserved reserved reserved reserved 74 | CATALOG non-reserved non-reserved non-reserved reserved 75 | CATALOG_NAME non-reserved non-reserved non-reserved 76 | CEIL reserved reserved 77 | CEILING reserved reserved 78 | CHAIN non-reserved non-reserved non-reserved 79 | CHAINING non-reserved non-reserved 80 | CHAR non-reserved (cannot be function or type), requires AS reserved reserved reserved 81 | CHARACTER non-reserved (cannot be function or type), requires AS reserved reserved reserved 82 | CHARACTERISTICS non-reserved non-reserved non-reserved 83 | CHARACTERS non-reserved non-reserved 84 | CHARACTER_LENGTH reserved reserved reserved 85 | CHARACTER_​SET_​CATALOG non-reserved non-reserved non-reserved 86 | CHARACTER_SET_NAME non-reserved non-reserved non-reserved 87 | CHARACTER_SET_SCHEMA non-reserved non-reserved non-reserved 88 | CHAR_LENGTH reserved reserved reserved 89 | CHECK reserved reserved reserved reserved 90 | CHECKPOINT non-reserved 91 | CLASS non-reserved 92 | CLASSIFIER reserved reserved 93 | CLASS_ORIGIN non-reserved non-reserved non-reserved 94 | CLOB reserved reserved 95 | CLOSE non-reserved reserved reserved reserved 96 | CLUSTER non-reserved 97 | COALESCE non-reserved (cannot be function or type) reserved reserved reserved 98 | COBOL non-reserved non-reserved non-reserved 99 | COLLATE reserved reserved reserved reserved 100 | COLLATION reserved (can be function or type) non-reserved non-reserved reserved 101 | COLLATION_CATALOG non-reserved non-reserved non-reserved 102 | COLLATION_NAME non-reserved non-reserved non-reserved 103 | COLLATION_SCHEMA non-reserved non-reserved non-reserved 104 | COLLECT reserved reserved 105 | COLUMN reserved reserved reserved reserved 106 | COLUMNS non-reserved non-reserved non-reserved 107 | COLUMN_NAME non-reserved non-reserved non-reserved 108 | COMMAND_FUNCTION non-reserved non-reserved non-reserved 109 | COMMAND_​FUNCTION_​CODE non-reserved non-reserved 110 | COMMENT non-reserved 111 | COMMENTS non-reserved 112 | COMMIT non-reserved reserved reserved reserved 113 | COMMITTED non-reserved non-reserved non-reserved non-reserved 114 | COMPRESSION non-reserved 115 | CONCURRENTLY reserved (can be function or type) 116 | CONDITION reserved reserved 117 | CONDITIONAL non-reserved non-reserved 118 | CONDITION_NUMBER non-reserved non-reserved non-reserved 119 | CONFIGURATION non-reserved 120 | CONFLICT non-reserved 121 | CONNECT reserved reserved reserved 122 | CONNECTION non-reserved non-reserved non-reserved reserved 123 | CONNECTION_NAME non-reserved non-reserved non-reserved 124 | CONSTRAINT reserved reserved reserved reserved 125 | CONSTRAINTS non-reserved non-reserved non-reserved reserved 126 | CONSTRAINT_CATALOG non-reserved non-reserved non-reserved 127 | CONSTRAINT_NAME non-reserved non-reserved non-reserved 128 | CONSTRAINT_SCHEMA non-reserved non-reserved non-reserved 129 | CONSTRUCTOR non-reserved non-reserved 130 | CONTAINS reserved reserved 131 | CONTENT non-reserved non-reserved non-reserved 132 | CONTINUE non-reserved non-reserved non-reserved reserved 133 | CONTROL non-reserved non-reserved 134 | CONVERSION non-reserved 135 | CONVERT reserved reserved reserved 136 | COPARTITION non-reserved 137 | COPY non-reserved reserved reserved 138 | CORR reserved reserved 139 | CORRESPONDING reserved reserved reserved 140 | COS reserved reserved 141 | COSH reserved reserved 142 | COST non-reserved 143 | COUNT reserved reserved reserved 144 | COVAR_POP reserved reserved 145 | COVAR_SAMP reserved reserved 146 | CREATE reserved, requires AS reserved reserved reserved 147 | CROSS reserved (can be function or type) reserved reserved reserved 148 | CSV non-reserved 149 | CUBE non-reserved reserved reserved 150 | CUME_DIST reserved reserved 151 | CURRENT non-reserved reserved reserved reserved 152 | CURRENT_CATALOG reserved reserved reserved 153 | CURRENT_DATE reserved reserved reserved reserved 154 | CURRENT_​DEFAULT_​TRANSFORM_​GROUP reserved reserved 155 | CURRENT_PATH reserved reserved 156 | CURRENT_ROLE reserved reserved reserved 157 | CURRENT_ROW reserved reserved 158 | CURRENT_SCHEMA reserved (can be function or type) reserved reserved 159 | CURRENT_TIME reserved reserved reserved reserved 160 | CURRENT_TIMESTAMP reserved reserved reserved reserved 161 | CURRENT_​TRANSFORM_​GROUP_​FOR_​TYPE reserved reserved 162 | CURRENT_USER reserved reserved reserved reserved 163 | CURSOR non-reserved reserved reserved reserved 164 | CURSOR_NAME non-reserved non-reserved non-reserved 165 | CYCLE non-reserved reserved reserved 166 | DATA non-reserved non-reserved non-reserved non-reserved 167 | DATABASE non-reserved 168 | DATALINK reserved reserved 169 | DATE reserved reserved reserved 170 | DATETIME_​INTERVAL_​CODE non-reserved non-reserved non-reserved 171 | DATETIME_​INTERVAL_​PRECISION non-reserved non-reserved non-reserved 172 | DAY non-reserved, requires AS reserved reserved reserved 173 | DB non-reserved non-reserved 174 | DEALLOCATE non-reserved reserved reserved reserved 175 | DEC non-reserved (cannot be function or type) reserved reserved reserved 176 | DECFLOAT reserved reserved 177 | DECIMAL non-reserved (cannot be function or type) reserved reserved reserved 178 | DECLARE non-reserved reserved reserved reserved 179 | DEFAULT reserved reserved reserved reserved 180 | DEFAULTS non-reserved non-reserved non-reserved 181 | DEFERRABLE reserved non-reserved non-reserved reserved 182 | DEFERRED non-reserved non-reserved non-reserved reserved 183 | DEFINE reserved reserved 184 | DEFINED non-reserved non-reserved 185 | DEFINER non-reserved non-reserved non-reserved 186 | DEGREE non-reserved non-reserved 187 | DELETE non-reserved reserved reserved reserved 188 | DELIMITER non-reserved 189 | DELIMITERS non-reserved 190 | DENSE_RANK reserved reserved 191 | DEPENDS non-reserved 192 | DEPTH non-reserved non-reserved non-reserved 193 | DEREF reserved reserved 194 | DERIVED non-reserved non-reserved 195 | DESC reserved non-reserved non-reserved reserved 196 | DESCRIBE reserved reserved reserved 197 | DESCRIPTOR non-reserved non-reserved reserved 198 | DETACH non-reserved 199 | DETERMINISTIC reserved reserved 200 | DIAGNOSTICS non-reserved non-reserved reserved 201 | DICTIONARY non-reserved 202 | DISABLE non-reserved 203 | DISCARD non-reserved 204 | DISCONNECT reserved reserved reserved 205 | DISPATCH non-reserved non-reserved 206 | DISTINCT reserved reserved reserved reserved 207 | DLNEWCOPY reserved reserved 208 | DLPREVIOUSCOPY reserved reserved 209 | DLURLCOMPLETE reserved reserved 210 | DLURLCOMPLETEONLY reserved reserved 211 | DLURLCOMPLETEWRITE reserved reserved 212 | DLURLPATH reserved reserved 213 | DLURLPATHONLY reserved reserved 214 | DLURLPATHWRITE reserved reserved 215 | DLURLSCHEME reserved reserved 216 | DLURLSERVER reserved reserved 217 | DLVALUE reserved reserved 218 | DO reserved 219 | DOCUMENT non-reserved non-reserved non-reserved 220 | DOMAIN non-reserved non-reserved non-reserved reserved 221 | DOUBLE non-reserved reserved reserved reserved 222 | DROP non-reserved reserved reserved reserved 223 | DYNAMIC reserved reserved 224 | DYNAMIC_FUNCTION non-reserved non-reserved non-reserved 225 | DYNAMIC_​FUNCTION_​CODE non-reserved non-reserved 226 | EACH non-reserved reserved reserved 227 | ELEMENT reserved reserved 228 | ELSE reserved reserved reserved reserved 229 | EMPTY reserved reserved 230 | ENABLE non-reserved 231 | ENCODING non-reserved non-reserved non-reserved 232 | ENCRYPTED non-reserved 233 | END reserved reserved reserved reserved 234 | END-EXEC reserved reserved reserved 235 | END_FRAME reserved reserved 236 | END_PARTITION reserved reserved 237 | ENFORCED non-reserved non-reserved 238 | ENUM non-reserved 239 | EQUALS reserved reserved 240 | ERROR non-reserved non-reserved 241 | ESCAPE non-reserved reserved reserved reserved 242 | EVENT non-reserved 243 | EVERY reserved reserved 244 | EXCEPT reserved, requires AS reserved reserved reserved 245 | EXCEPTION reserved 246 | EXCLUDE non-reserved non-reserved non-reserved 247 | EXCLUDING non-reserved non-reserved non-reserved 248 | EXCLUSIVE non-reserved 249 | EXEC reserved reserved reserved 250 | EXECUTE non-reserved reserved reserved reserved 251 | EXISTS non-reserved (cannot be function or type) reserved reserved reserved 252 | EXP reserved reserved 253 | EXPLAIN non-reserved 254 | EXPRESSION non-reserved non-reserved non-reserved 255 | EXTENSION non-reserved 256 | EXTERNAL non-reserved reserved reserved reserved 257 | EXTRACT non-reserved (cannot be function or type) reserved reserved reserved 258 | FALSE reserved reserved reserved reserved 259 | FAMILY non-reserved 260 | FETCH reserved, requires AS reserved reserved reserved 261 | FILE non-reserved non-reserved 262 | FILTER non-reserved, requires AS reserved reserved 263 | FINAL non-reserved non-reserved 264 | FINALIZE non-reserved 265 | FINISH non-reserved non-reserved 266 | FIRST non-reserved non-reserved non-reserved reserved 267 | FIRST_VALUE reserved reserved 268 | FLAG non-reserved non-reserved 269 | FLOAT non-reserved (cannot be function or type) reserved reserved reserved 270 | FLOOR reserved reserved 271 | FOLLOWING non-reserved non-reserved non-reserved 272 | FOR reserved, requires AS reserved reserved reserved 273 | FORCE non-reserved 274 | FOREIGN reserved reserved reserved reserved 275 | FORMAT non-reserved non-reserved non-reserved 276 | FORTRAN non-reserved non-reserved non-reserved 277 | FORWARD non-reserved 278 | FOUND non-reserved non-reserved reserved 279 | FRAME_ROW reserved reserved 280 | FREE reserved reserved 281 | FREEZE reserved (can be function or type) 282 | FROM reserved, requires AS reserved reserved reserved 283 | FS non-reserved non-reserved 284 | FULFILL non-reserved non-reserved 285 | FULL reserved (can be function or type) reserved reserved reserved 286 | FUNCTION non-reserved reserved reserved 287 | FUNCTIONS non-reserved 288 | FUSION reserved reserved 289 | G non-reserved non-reserved 290 | GENERAL non-reserved non-reserved 291 | GENERATED non-reserved non-reserved non-reserved 292 | GET reserved reserved reserved 293 | GLOBAL non-reserved reserved reserved reserved 294 | GO non-reserved non-reserved reserved 295 | GOTO non-reserved non-reserved reserved 296 | GRANT reserved, requires AS reserved reserved reserved 297 | GRANTED non-reserved non-reserved non-reserved 298 | GREATEST non-reserved (cannot be function or type) reserved 299 | GROUP reserved, requires AS reserved reserved reserved 300 | GROUPING non-reserved (cannot be function or type) reserved reserved 301 | GROUPS non-reserved reserved reserved 302 | HANDLER non-reserved 303 | HAVING reserved, requires AS reserved reserved reserved 304 | HEADER non-reserved 305 | HEX non-reserved non-reserved 306 | HIERARCHY non-reserved non-reserved 307 | HOLD non-reserved reserved reserved 308 | HOUR non-reserved, requires AS reserved reserved reserved 309 | ID non-reserved non-reserved 310 | IDENTITY non-reserved reserved reserved reserved 311 | IF non-reserved 312 | IGNORE non-reserved non-reserved 313 | ILIKE reserved (can be function or type) 314 | IMMEDIATE non-reserved non-reserved non-reserved reserved 315 | IMMEDIATELY non-reserved non-reserved 316 | IMMUTABLE non-reserved 317 | IMPLEMENTATION non-reserved non-reserved 318 | IMPLICIT non-reserved 319 | IMPORT non-reserved reserved reserved 320 | IN reserved reserved reserved reserved 321 | INCLUDE non-reserved 322 | INCLUDING non-reserved non-reserved non-reserved 323 | INCREMENT non-reserved non-reserved non-reserved 324 | INDENT non-reserved non-reserved non-reserved 325 | INDEX non-reserved 326 | INDEXES non-reserved 327 | INDICATOR reserved reserved reserved 328 | INHERIT non-reserved 329 | INHERITS non-reserved 330 | INITIAL reserved reserved 331 | INITIALLY reserved non-reserved non-reserved reserved 332 | INLINE non-reserved 333 | INNER reserved (can be function or type) reserved reserved reserved 334 | INOUT non-reserved (cannot be function or type) reserved reserved 335 | INPUT non-reserved non-reserved non-reserved reserved 336 | INSENSITIVE non-reserved reserved reserved reserved 337 | INSERT non-reserved reserved reserved reserved 338 | INSTANCE non-reserved non-reserved 339 | INSTANTIABLE non-reserved non-reserved 340 | INSTEAD non-reserved non-reserved non-reserved 341 | INT non-reserved (cannot be function or type) reserved reserved reserved 342 | INTEGER non-reserved (cannot be function or type) reserved reserved reserved 343 | INTEGRITY non-reserved non-reserved 344 | INTERSECT reserved, requires AS reserved reserved reserved 345 | INTERSECTION reserved reserved 346 | INTERVAL non-reserved (cannot be function or type) reserved reserved reserved 347 | INTO reserved, requires AS reserved reserved reserved 348 | INVOKER non-reserved non-reserved non-reserved 349 | IS reserved (can be function or type) reserved reserved reserved 350 | ISNULL reserved (can be function or type), requires AS 351 | ISOLATION non-reserved non-reserved non-reserved reserved 352 | JOIN reserved (can be function or type) reserved reserved reserved 353 | JSON non-reserved reserved 354 | JSON_ARRAY non-reserved (cannot be function or type) reserved reserved 355 | JSON_ARRAYAGG non-reserved (cannot be function or type) reserved reserved 356 | JSON_EXISTS reserved reserved 357 | JSON_OBJECT non-reserved (cannot be function or type) reserved reserved 358 | JSON_OBJECTAGG non-reserved (cannot be function or type) reserved reserved 359 | JSON_QUERY reserved reserved 360 | JSON_SCALAR reserved 361 | JSON_SERIALIZE reserved 362 | JSON_TABLE reserved reserved 363 | JSON_TABLE_PRIMITIVE reserved reserved 364 | JSON_VALUE reserved reserved 365 | K non-reserved non-reserved 366 | KEEP non-reserved non-reserved 367 | KEY non-reserved non-reserved non-reserved reserved 368 | KEYS non-reserved non-reserved non-reserved 369 | KEY_MEMBER non-reserved non-reserved 370 | KEY_TYPE non-reserved non-reserved 371 | LABEL non-reserved 372 | LAG reserved reserved 373 | LANGUAGE non-reserved reserved reserved reserved 374 | LARGE non-reserved reserved reserved 375 | LAST non-reserved non-reserved non-reserved reserved 376 | LAST_VALUE reserved reserved 377 | LATERAL reserved reserved reserved 378 | LEAD reserved reserved 379 | LEADING reserved reserved reserved reserved 380 | LEAKPROOF non-reserved 381 | LEAST non-reserved (cannot be function or type) reserved 382 | LEFT reserved (can be function or type) reserved reserved reserved 383 | LENGTH non-reserved non-reserved non-reserved 384 | LEVEL non-reserved non-reserved non-reserved reserved 385 | LIBRARY non-reserved non-reserved 386 | LIKE reserved (can be function or type) reserved reserved reserved 387 | LIKE_REGEX reserved reserved 388 | LIMIT reserved, requires AS non-reserved non-reserved 389 | LINK non-reserved non-reserved 390 | LISTAGG reserved reserved 391 | LISTEN non-reserved 392 | LN reserved reserved 393 | LOAD non-reserved 394 | LOCAL non-reserved reserved reserved reserved 395 | LOCALTIME reserved reserved reserved 396 | LOCALTIMESTAMP reserved reserved reserved 397 | LOCATION non-reserved non-reserved non-reserved 398 | LOCATOR non-reserved non-reserved 399 | LOCK non-reserved 400 | LOCKED non-reserved 401 | LOG reserved reserved 402 | LOG10 reserved reserved 403 | LOGGED non-reserved 404 | LOWER reserved reserved reserved 405 | LPAD reserved 406 | LTRIM reserved 407 | M non-reserved non-reserved 408 | MAP non-reserved non-reserved 409 | MAPPING non-reserved non-reserved non-reserved 410 | MATCH non-reserved reserved reserved reserved 411 | MATCHED non-reserved non-reserved non-reserved 412 | MATCHES reserved reserved 413 | MATCH_NUMBER reserved reserved 414 | MATCH_RECOGNIZE reserved reserved 415 | MATERIALIZED non-reserved 416 | MAX reserved reserved reserved 417 | MAXVALUE non-reserved non-reserved non-reserved 418 | MEASURES non-reserved non-reserved 419 | MEMBER reserved reserved 420 | MERGE non-reserved reserved reserved 421 | MESSAGE_LENGTH non-reserved non-reserved non-reserved 422 | MESSAGE_OCTET_LENGTH non-reserved non-reserved non-reserved 423 | MESSAGE_TEXT non-reserved non-reserved non-reserved 424 | METHOD non-reserved reserved reserved 425 | MIN reserved reserved reserved 426 | MINUTE non-reserved, requires AS reserved reserved reserved 427 | MINVALUE non-reserved non-reserved non-reserved 428 | MOD reserved reserved 429 | MODE non-reserved 430 | MODIFIES reserved reserved 431 | MODULE reserved reserved reserved 432 | MONTH non-reserved, requires AS reserved reserved reserved 433 | MORE non-reserved non-reserved non-reserved 434 | MOVE non-reserved 435 | MULTISET reserved reserved 436 | MUMPS non-reserved non-reserved non-reserved 437 | NAME non-reserved non-reserved non-reserved non-reserved 438 | NAMES non-reserved non-reserved non-reserved reserved 439 | NAMESPACE non-reserved non-reserved 440 | NATIONAL non-reserved (cannot be function or type) reserved reserved reserved 441 | NATURAL reserved (can be function or type) reserved reserved reserved 442 | NCHAR non-reserved (cannot be function or type) reserved reserved reserved 443 | NCLOB reserved reserved 444 | NESTED non-reserved non-reserved 445 | NESTING non-reserved non-reserved 446 | NEW non-reserved reserved reserved 447 | NEXT non-reserved non-reserved non-reserved reserved 448 | NFC non-reserved non-reserved non-reserved 449 | NFD non-reserved non-reserved non-reserved 450 | NFKC non-reserved non-reserved non-reserved 451 | NFKD non-reserved non-reserved non-reserved 452 | NIL non-reserved non-reserved 453 | NO non-reserved reserved reserved reserved 454 | NONE non-reserved (cannot be function or type) reserved reserved 455 | NORMALIZE non-reserved (cannot be function or type) reserved reserved 456 | NORMALIZED non-reserved non-reserved non-reserved 457 | NOT reserved reserved reserved reserved 458 | NOTHING non-reserved 459 | NOTIFY non-reserved 460 | NOTNULL reserved (can be function or type), requires AS 461 | NOWAIT non-reserved 462 | NTH_VALUE reserved reserved 463 | NTILE reserved reserved 464 | NULL reserved reserved reserved reserved 465 | NULLABLE non-reserved non-reserved non-reserved 466 | NULLIF non-reserved (cannot be function or type) reserved reserved reserved 467 | NULLS non-reserved non-reserved non-reserved 468 | NULL_ORDERING non-reserved non-reserved 469 | NUMBER non-reserved non-reserved non-reserved 470 | NUMERIC non-reserved (cannot be function or type) reserved reserved reserved 471 | OBJECT non-reserved non-reserved non-reserved 472 | OCCURRENCE non-reserved non-reserved 473 | OCCURRENCES_REGEX reserved reserved 474 | OCTETS non-reserved non-reserved 475 | OCTET_LENGTH reserved reserved reserved 476 | OF non-reserved reserved reserved reserved 477 | OFF non-reserved non-reserved non-reserved 478 | OFFSET reserved, requires AS reserved reserved 479 | OIDS non-reserved 480 | OLD non-reserved reserved reserved 481 | OMIT reserved reserved 482 | ON reserved, requires AS reserved reserved reserved 483 | ONE reserved reserved 484 | ONLY reserved reserved reserved reserved 485 | OPEN reserved reserved reserved 486 | OPERATOR non-reserved 487 | OPTION non-reserved non-reserved non-reserved reserved 488 | OPTIONS non-reserved non-reserved non-reserved 489 | OR reserved reserved reserved reserved 490 | ORDER reserved, requires AS reserved reserved reserved 491 | ORDERING non-reserved non-reserved 492 | ORDINALITY non-reserved non-reserved non-reserved 493 | OTHERS non-reserved non-reserved non-reserved 494 | OUT non-reserved (cannot be function or type) reserved reserved 495 | OUTER reserved (can be function or type) reserved reserved reserved 496 | OUTPUT non-reserved non-reserved reserved 497 | OVER non-reserved, requires AS reserved reserved 498 | OVERFLOW non-reserved non-reserved 499 | OVERLAPS reserved (can be function or type), requires AS reserved reserved reserved 500 | OVERLAY non-reserved (cannot be function or type) reserved reserved 501 | OVERRIDING non-reserved non-reserved non-reserved 502 | OWNED non-reserved 503 | OWNER non-reserved 504 | P non-reserved non-reserved 505 | PAD non-reserved non-reserved reserved 506 | PARALLEL non-reserved 507 | PARAMETER non-reserved reserved reserved 508 | PARAMETER_MODE non-reserved non-reserved 509 | PARAMETER_NAME non-reserved non-reserved 510 | PARAMETER_​ORDINAL_​POSITION non-reserved non-reserved 511 | PARAMETER_​SPECIFIC_​CATALOG non-reserved non-reserved 512 | PARAMETER_​SPECIFIC_​NAME non-reserved non-reserved 513 | PARAMETER_​SPECIFIC_​SCHEMA non-reserved non-reserved 514 | PARSER non-reserved 515 | PARTIAL non-reserved non-reserved non-reserved reserved 516 | PARTITION non-reserved reserved reserved 517 | PASCAL non-reserved non-reserved non-reserved 518 | PASS non-reserved non-reserved 519 | PASSING non-reserved non-reserved non-reserved 520 | PASSTHROUGH non-reserved non-reserved 521 | PASSWORD non-reserved 522 | PAST non-reserved non-reserved 523 | PATH non-reserved non-reserved 524 | PATTERN reserved reserved 525 | PER reserved reserved 526 | PERCENT reserved reserved 527 | PERCENTILE_CONT reserved reserved 528 | PERCENTILE_DISC reserved reserved 529 | PERCENT_RANK reserved reserved 530 | PERIOD reserved reserved 531 | PERMISSION non-reserved non-reserved 532 | PERMUTE non-reserved non-reserved 533 | PIPE non-reserved non-reserved 534 | PLACING reserved non-reserved non-reserved 535 | PLAN non-reserved non-reserved 536 | PLANS non-reserved 537 | PLI non-reserved non-reserved non-reserved 538 | POLICY non-reserved 539 | PORTION reserved reserved 540 | POSITION non-reserved (cannot be function or type) reserved reserved reserved 541 | POSITION_REGEX reserved reserved 542 | POWER reserved reserved 543 | PRECEDES reserved reserved 544 | PRECEDING non-reserved non-reserved non-reserved 545 | PRECISION non-reserved (cannot be function or type), requires AS reserved reserved reserved 546 | PREPARE non-reserved reserved reserved reserved 547 | PREPARED non-reserved 548 | PRESERVE non-reserved non-reserved non-reserved reserved 549 | PREV non-reserved non-reserved 550 | PRIMARY reserved reserved reserved reserved 551 | PRIOR non-reserved non-reserved non-reserved reserved 552 | PRIVATE non-reserved non-reserved 553 | PRIVILEGES non-reserved non-reserved non-reserved reserved 554 | PROCEDURAL non-reserved 555 | PROCEDURE non-reserved reserved reserved reserved 556 | PROCEDURES non-reserved 557 | PROGRAM non-reserved 558 | PRUNE non-reserved non-reserved 559 | PTF reserved reserved 560 | PUBLIC non-reserved non-reserved reserved 561 | PUBLICATION non-reserved 562 | QUOTE non-reserved 563 | QUOTES non-reserved non-reserved 564 | RANGE non-reserved reserved reserved 565 | RANK reserved reserved 566 | READ non-reserved non-reserved non-reserved reserved 567 | READS reserved reserved 568 | REAL non-reserved (cannot be function or type) reserved reserved reserved 569 | REASSIGN non-reserved 570 | RECHECK non-reserved 571 | RECOVERY non-reserved non-reserved 572 | RECURSIVE non-reserved reserved reserved 573 | REF non-reserved reserved reserved 574 | REFERENCES reserved reserved reserved reserved 575 | REFERENCING non-reserved reserved reserved 576 | REFRESH non-reserved 577 | REGR_AVGX reserved reserved 578 | REGR_AVGY reserved reserved 579 | REGR_COUNT reserved reserved 580 | REGR_INTERCEPT reserved reserved 581 | REGR_R2 reserved reserved 582 | REGR_SLOPE reserved reserved 583 | REGR_SXX reserved reserved 584 | REGR_SXY reserved reserved 585 | REGR_SYY reserved reserved 586 | REINDEX non-reserved 587 | RELATIVE non-reserved non-reserved non-reserved reserved 588 | RELEASE non-reserved reserved reserved 589 | RENAME non-reserved 590 | REPEATABLE non-reserved non-reserved non-reserved non-reserved 591 | REPLACE non-reserved 592 | REPLICA non-reserved 593 | REQUIRING non-reserved non-reserved 594 | RESET non-reserved 595 | RESPECT non-reserved non-reserved 596 | RESTART non-reserved non-reserved non-reserved 597 | RESTORE non-reserved non-reserved 598 | RESTRICT non-reserved non-reserved non-reserved reserved 599 | RESULT reserved reserved 600 | RETURN non-reserved reserved reserved 601 | RETURNED_CARDINALITY non-reserved non-reserved 602 | RETURNED_LENGTH non-reserved non-reserved non-reserved 603 | RETURNED_​OCTET_​LENGTH non-reserved non-reserved non-reserved 604 | RETURNED_SQLSTATE non-reserved non-reserved non-reserved 605 | RETURNING reserved, requires AS non-reserved non-reserved 606 | RETURNS non-reserved reserved reserved 607 | REVOKE non-reserved reserved reserved reserved 608 | RIGHT reserved (can be function or type) reserved reserved reserved 609 | ROLE non-reserved non-reserved non-reserved 610 | ROLLBACK non-reserved reserved reserved reserved 611 | ROLLUP non-reserved reserved reserved 612 | ROUTINE non-reserved non-reserved non-reserved 613 | ROUTINES non-reserved 614 | ROUTINE_CATALOG non-reserved non-reserved 615 | ROUTINE_NAME non-reserved non-reserved 616 | ROUTINE_SCHEMA non-reserved non-reserved 617 | ROW non-reserved (cannot be function or type) reserved reserved 618 | ROWS non-reserved reserved reserved reserved 619 | ROW_COUNT non-reserved non-reserved non-reserved 620 | ROW_NUMBER reserved reserved 621 | RPAD reserved 622 | RTRIM reserved 623 | RULE non-reserved 624 | RUNNING reserved reserved 625 | SAVEPOINT non-reserved reserved reserved 626 | SCALAR non-reserved non-reserved non-reserved 627 | SCALE non-reserved non-reserved non-reserved 628 | SCHEMA non-reserved non-reserved non-reserved reserved 629 | SCHEMAS non-reserved 630 | SCHEMA_NAME non-reserved non-reserved non-reserved 631 | SCOPE reserved reserved 632 | SCOPE_CATALOG non-reserved non-reserved 633 | SCOPE_NAME non-reserved non-reserved 634 | SCOPE_SCHEMA non-reserved non-reserved 635 | SCROLL non-reserved reserved reserved reserved 636 | SEARCH non-reserved reserved reserved 637 | SECOND non-reserved, requires AS reserved reserved reserved 638 | SECTION non-reserved non-reserved reserved 639 | SECURITY non-reserved non-reserved non-reserved 640 | SEEK reserved reserved 641 | SELECT reserved reserved reserved reserved 642 | SELECTIVE non-reserved non-reserved 643 | SELF non-reserved non-reserved 644 | SEMANTICS non-reserved non-reserved 645 | SENSITIVE reserved reserved 646 | SEQUENCE non-reserved non-reserved non-reserved 647 | SEQUENCES non-reserved 648 | SERIALIZABLE non-reserved non-reserved non-reserved non-reserved 649 | SERVER non-reserved non-reserved non-reserved 650 | SERVER_NAME non-reserved non-reserved non-reserved 651 | SESSION non-reserved non-reserved non-reserved reserved 652 | SESSION_USER reserved reserved reserved reserved 653 | SET non-reserved reserved reserved reserved 654 | SETOF non-reserved (cannot be function or type) 655 | SETS non-reserved non-reserved non-reserved 656 | SHARE non-reserved 657 | SHOW non-reserved reserved reserved 658 | SIMILAR reserved (can be function or type) reserved reserved 659 | SIMPLE non-reserved non-reserved non-reserved 660 | SIN reserved reserved 661 | SINH reserved reserved 662 | SIZE non-reserved non-reserved reserved 663 | SKIP non-reserved reserved reserved 664 | SMALLINT non-reserved (cannot be function or type) reserved reserved reserved 665 | SNAPSHOT non-reserved 666 | SOME reserved reserved reserved reserved 667 | SORT_DIRECTION non-reserved non-reserved 668 | SOURCE non-reserved non-reserved 669 | SPACE non-reserved non-reserved reserved 670 | SPECIFIC reserved reserved 671 | SPECIFICTYPE reserved reserved 672 | SPECIFIC_NAME non-reserved non-reserved 673 | SQL non-reserved reserved reserved reserved 674 | SQLCODE reserved 675 | SQLERROR reserved 676 | SQLEXCEPTION reserved reserved 677 | SQLSTATE reserved reserved reserved 678 | SQLWARNING reserved reserved 679 | SQRT reserved reserved 680 | STABLE non-reserved 681 | STANDALONE non-reserved non-reserved non-reserved 682 | START non-reserved reserved reserved 683 | STATE non-reserved non-reserved 684 | STATEMENT non-reserved non-reserved non-reserved 685 | STATIC reserved reserved 686 | STATISTICS non-reserved 687 | STDDEV_POP reserved reserved 688 | STDDEV_SAMP reserved reserved 689 | STDIN non-reserved 690 | STDOUT non-reserved 691 | STORAGE non-reserved 692 | STORED non-reserved 693 | STRICT non-reserved 694 | STRING non-reserved non-reserved 695 | STRIP non-reserved non-reserved non-reserved 696 | STRUCTURE non-reserved non-reserved 697 | STYLE non-reserved non-reserved 698 | SUBCLASS_ORIGIN non-reserved non-reserved non-reserved 699 | SUBMULTISET reserved reserved 700 | SUBSCRIPTION non-reserved 701 | SUBSET reserved reserved 702 | SUBSTRING non-reserved (cannot be function or type) reserved reserved reserved 703 | SUBSTRING_REGEX reserved reserved 704 | SUCCEEDS reserved reserved 705 | SUM reserved reserved reserved 706 | SUPPORT non-reserved 707 | SYMMETRIC reserved reserved reserved 708 | SYSID non-reserved 709 | SYSTEM non-reserved reserved reserved 710 | SYSTEM_TIME reserved reserved 711 | SYSTEM_USER reserved reserved reserved reserved 712 | T non-reserved non-reserved 713 | TABLE reserved reserved reserved reserved 714 | TABLES non-reserved 715 | TABLESAMPLE reserved (can be function or type) reserved reserved 716 | TABLESPACE non-reserved 717 | TABLE_NAME non-reserved non-reserved non-reserved 718 | TAN reserved reserved 719 | TANH reserved reserved 720 | TEMP non-reserved 721 | TEMPLATE non-reserved 722 | TEMPORARY non-reserved non-reserved non-reserved reserved 723 | TEXT non-reserved 724 | THEN reserved reserved reserved reserved 725 | THROUGH non-reserved non-reserved 726 | TIES non-reserved non-reserved non-reserved 727 | TIME non-reserved (cannot be function or type) reserved reserved reserved 728 | TIMESTAMP non-reserved (cannot be function or type) reserved reserved reserved 729 | TIMEZONE_HOUR reserved reserved reserved 730 | TIMEZONE_MINUTE reserved reserved reserved 731 | TO reserved, requires AS reserved reserved reserved 732 | TOKEN non-reserved non-reserved 733 | TOP_LEVEL_COUNT non-reserved non-reserved 734 | TRAILING reserved reserved reserved reserved 735 | TRANSACTION non-reserved non-reserved non-reserved reserved 736 | TRANSACTIONS_​COMMITTED non-reserved non-reserved 737 | TRANSACTIONS_​ROLLED_​BACK non-reserved non-reserved 738 | TRANSACTION_ACTIVE non-reserved non-reserved 739 | TRANSFORM non-reserved non-reserved non-reserved 740 | TRANSFORMS non-reserved non-reserved 741 | TRANSLATE reserved reserved reserved 742 | TRANSLATE_REGEX reserved reserved 743 | TRANSLATION reserved reserved reserved 744 | TREAT non-reserved (cannot be function or type) reserved reserved 745 | TRIGGER non-reserved reserved reserved 746 | TRIGGER_CATALOG non-reserved non-reserved 747 | TRIGGER_NAME non-reserved non-reserved 748 | TRIGGER_SCHEMA non-reserved non-reserved 749 | TRIM non-reserved (cannot be function or type) reserved reserved reserved 750 | TRIM_ARRAY reserved reserved 751 | TRUE reserved reserved reserved reserved 752 | TRUNCATE non-reserved reserved reserved 753 | TRUSTED non-reserved 754 | TYPE non-reserved non-reserved non-reserved non-reserved 755 | TYPES non-reserved 756 | UESCAPE non-reserved reserved reserved 757 | UNBOUNDED non-reserved non-reserved non-reserved 758 | UNCOMMITTED non-reserved non-reserved non-reserved non-reserved 759 | UNCONDITIONAL non-reserved non-reserved 760 | UNDER non-reserved non-reserved 761 | UNENCRYPTED non-reserved 762 | UNION reserved, requires AS reserved reserved reserved 763 | UNIQUE reserved reserved reserved reserved 764 | UNKNOWN non-reserved reserved reserved reserved 765 | UNLINK non-reserved non-reserved 766 | UNLISTEN non-reserved 767 | UNLOGGED non-reserved 768 | UNMATCHED non-reserved non-reserved 769 | UNNAMED non-reserved non-reserved non-reserved 770 | UNNEST reserved reserved 771 | UNTIL non-reserved 772 | UNTYPED non-reserved non-reserved 773 | UPDATE non-reserved reserved reserved reserved 774 | UPPER reserved reserved reserved 775 | URI non-reserved non-reserved 776 | USAGE non-reserved non-reserved reserved 777 | USER reserved reserved reserved reserved 778 | USER_​DEFINED_​TYPE_​CATALOG non-reserved non-reserved 779 | USER_​DEFINED_​TYPE_​CODE non-reserved non-reserved 780 | USER_​DEFINED_​TYPE_​NAME non-reserved non-reserved 781 | USER_​DEFINED_​TYPE_​SCHEMA non-reserved non-reserved 782 | USING reserved reserved reserved reserved 783 | UTF16 non-reserved non-reserved 784 | UTF32 non-reserved non-reserved 785 | UTF8 non-reserved non-reserved 786 | VACUUM non-reserved 787 | VALID non-reserved non-reserved non-reserved 788 | VALIDATE non-reserved 789 | VALIDATOR non-reserved 790 | VALUE non-reserved reserved reserved reserved 791 | VALUES non-reserved (cannot be function or type) reserved reserved reserved 792 | VALUE_OF reserved reserved 793 | VARBINARY reserved reserved 794 | VARCHAR non-reserved (cannot be function or type) reserved reserved reserved 795 | VARIADIC reserved 796 | VARYING non-reserved, requires AS reserved reserved reserved 797 | VAR_POP reserved reserved 798 | VAR_SAMP reserved reserved 799 | VERBOSE reserved (can be function or type) 800 | VERSION non-reserved non-reserved non-reserved 801 | VERSIONING reserved reserved 802 | VIEW non-reserved non-reserved non-reserved reserved 803 | VIEWS non-reserved 804 | VOLATILE non-reserved 805 | WHEN reserved reserved reserved reserved 806 | WHENEVER reserved reserved reserved 807 | WHERE reserved, requires AS reserved reserved reserved 808 | WHITESPACE non-reserved non-reserved non-reserved 809 | WIDTH_BUCKET reserved reserved 810 | WINDOW reserved, requires AS reserved reserved 811 | WITH reserved, requires AS reserved reserved reserved 812 | WITHIN non-reserved, requires AS reserved reserved 813 | WITHOUT non-reserved, requires AS reserved reserved 814 | WORK non-reserved non-reserved non-reserved reserved 815 | WRAPPER non-reserved non-reserved non-reserved 816 | WRITE non-reserved non-reserved non-reserved reserved 817 | XML non-reserved reserved reserved 818 | XMLAGG reserved reserved 819 | XMLATTRIBUTES non-reserved (cannot be function or type) reserved reserved 820 | XMLBINARY reserved reserved 821 | XMLCAST reserved reserved 822 | XMLCOMMENT reserved reserved 823 | XMLCONCAT non-reserved (cannot be function or type) reserved reserved 824 | XMLDECLARATION non-reserved non-reserved 825 | XMLDOCUMENT reserved reserved 826 | XMLELEMENT non-reserved (cannot be function or type) reserved reserved 827 | XMLEXISTS non-reserved (cannot be function or type) reserved reserved 828 | XMLFOREST non-reserved (cannot be function or type) reserved reserved 829 | XMLITERATE reserved reserved 830 | XMLNAMESPACES non-reserved (cannot be function or type) reserved reserved 831 | XMLPARSE non-reserved (cannot be function or type) reserved reserved 832 | XMLPI non-reserved (cannot be function or type) reserved reserved 833 | XMLQUERY reserved reserved 834 | XMLROOT non-reserved (cannot be function or type) 835 | XMLSCHEMA non-reserved non-reserved 836 | XMLSERIALIZE non-reserved (cannot be function or type) reserved reserved 837 | XMLTABLE non-reserved (cannot be function or type) reserved reserved 838 | XMLTEXT reserved reserved 839 | XMLVALIDATE reserved reserved 840 | YEAR non-reserved, requires AS reserved reserved reserved 841 | YES non-reserved non-reserved non-reserved 842 | ZONE non-reserved non-reserved non-reserved reserved -------------------------------------------------------------------------------- /src/harlequin_postgres/loaders.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from datetime import date, datetime 4 | from typing import TYPE_CHECKING 5 | 6 | import psycopg 7 | from psycopg.errors import DataError 8 | 9 | # Subclass existing adapters so that the base case is handled normally. 10 | from psycopg.types.datetime import ( 11 | DateBinaryLoader, 12 | DateLoader, 13 | TimestampBinaryLoader, 14 | TimestampLoader, 15 | TimestamptzBinaryLoader, 16 | TimestamptzLoader, 17 | ) 18 | 19 | if TYPE_CHECKING: 20 | from psycopg.adapt import Buffer, Loader 21 | 22 | 23 | class InfDateLoader(DateLoader): 24 | def load(self, data: "Buffer") -> date: 25 | if data == b"infinity": 26 | return date.max 27 | elif data == b"-infinity": 28 | return date.min 29 | else: 30 | return super().load(data) 31 | 32 | 33 | class InfDateBinaryLoader(DateBinaryLoader): 34 | def load(self, data: "Buffer") -> date: 35 | try: 36 | return super().load(data) 37 | except DataError as e: 38 | if "date too small" in str(e): 39 | return date.min 40 | elif "date too large" in str(e): 41 | return date.max 42 | raise e 43 | 44 | 45 | class InfTimestampLoader(TimestampLoader): 46 | def load(self, data: "Buffer") -> datetime: 47 | if data == b"infinity": 48 | return datetime.max 49 | elif data == b"-infinity": 50 | return datetime.min 51 | else: 52 | return super().load(data) 53 | 54 | 55 | class InfTimestampBinaryLoader(TimestampBinaryLoader): 56 | def load(self, data: "Buffer") -> datetime: 57 | try: 58 | return super().load(data) 59 | except DataError as e: 60 | if "timestamp too small" in str(e): 61 | return datetime.min 62 | elif "timestamp too large" in str(e): 63 | return datetime.max 64 | raise e 65 | 66 | 67 | class InfTimestamptzLoader(TimestamptzLoader): 68 | def load(self, data: "Buffer") -> datetime: 69 | if data == b"infinity": 70 | return datetime.max 71 | elif data == b"-infinity": 72 | return datetime.min 73 | else: 74 | return super().load(data) 75 | 76 | 77 | class InfTimestamptzBinaryLoader(TimestamptzBinaryLoader): 78 | def load(self, data: "Buffer") -> datetime: 79 | try: 80 | return super().load(data) 81 | except DataError as e: 82 | if "timestamp too small" in str(e): 83 | return datetime.min 84 | elif "timestamp too large" in str(e): 85 | return datetime.max 86 | raise e 87 | 88 | 89 | INF_LOADERS: list[tuple[str, type["Loader"]]] = [ 90 | ("date", InfDateLoader), 91 | ("date", InfDateBinaryLoader), 92 | ("timestamp", InfTimestampLoader), 93 | ("timestamp", InfTimestampBinaryLoader), 94 | ("timestamptz", InfTimestamptzLoader), 95 | ("timestamptz", InfTimestamptzBinaryLoader), 96 | ] 97 | 98 | 99 | def register_inf_loaders() -> None: 100 | """ 101 | Register updated date/datetime loaders in the global types 102 | registry, so that any connections created afterwards will use 103 | the updated loaders (that allow infinity date/timestamps) 104 | """ 105 | for type_name, loader in INF_LOADERS: 106 | psycopg.adapters.register_loader(type_name, loader) 107 | -------------------------------------------------------------------------------- /src/harlequin_postgres/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tconbeer/harlequin-postgres/9d7c893723b98cd0bec9c8b0b89090135adc92d9/src/harlequin_postgres/py.typed -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | import sys 4 | from typing import Generator 5 | 6 | import psycopg 7 | import pytest 8 | from harlequin_postgres.adapter import ( 9 | HarlequinPostgresAdapter, 10 | HarlequinPostgresConnection, 11 | ) 12 | 13 | if sys.version_info < (3, 10): 14 | pass 15 | else: 16 | pass 17 | 18 | TEST_DB_CONN = "postgresql://postgres:for-testing@localhost:5432" 19 | 20 | 21 | @pytest.fixture 22 | def connection() -> Generator[HarlequinPostgresConnection, None, None]: 23 | pgconn = psycopg.connect(conninfo=TEST_DB_CONN, dbname="postgres") 24 | pgconn.autocommit = True 25 | cur = pgconn.cursor() 26 | cur.execute("drop database if exists test;") 27 | cur.execute("create database test;") 28 | cur.close() 29 | pgconn.close() 30 | conn = HarlequinPostgresAdapter( 31 | conn_str=(f"{TEST_DB_CONN}",), dbname="test" 32 | ).connect() 33 | yield conn 34 | conn.close() 35 | pgconn = psycopg.connect(conninfo=TEST_DB_CONN, dbname="postgres") 36 | pgconn.autocommit = True 37 | cur = pgconn.cursor() 38 | cur.execute("drop database if exists test;") 39 | cur.close() 40 | pgconn.close() 41 | -------------------------------------------------------------------------------- /tests/test_adapter.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | import sys 4 | from datetime import date, datetime 5 | 6 | import pytest 7 | from harlequin.adapter import HarlequinAdapter, HarlequinConnection, HarlequinCursor 8 | from harlequin.catalog import Catalog, CatalogItem 9 | from harlequin.exception import HarlequinConnectionError, HarlequinQueryError 10 | from harlequin_postgres.adapter import ( 11 | HarlequinPostgresAdapter, 12 | HarlequinPostgresConnection, 13 | ) 14 | from textual_fastdatatable.backend import create_backend 15 | 16 | if sys.version_info < (3, 10): 17 | from importlib_metadata import entry_points 18 | else: 19 | from importlib.metadata import entry_points 20 | 21 | TEST_DB_CONN = "postgresql://postgres:for-testing@localhost:5432" 22 | 23 | 24 | def test_plugin_discovery() -> None: 25 | PLUGIN_NAME = "postgres" 26 | eps = entry_points(group="harlequin.adapter") 27 | assert eps[PLUGIN_NAME] 28 | adapter_cls = eps[PLUGIN_NAME].load() 29 | assert issubclass(adapter_cls, HarlequinAdapter) 30 | assert adapter_cls == HarlequinPostgresAdapter 31 | 32 | 33 | def test_connect() -> None: 34 | conn = HarlequinPostgresAdapter(conn_str=(TEST_DB_CONN,)).connect() 35 | assert isinstance(conn, HarlequinConnection) 36 | 37 | 38 | def test_init_extra_kwargs() -> None: 39 | assert HarlequinPostgresAdapter( 40 | conn_str=(TEST_DB_CONN,), foo=1, bar="baz" 41 | ).connect() 42 | 43 | 44 | @pytest.mark.parametrize( 45 | "conn_str", 46 | [ 47 | ("foo",), 48 | ("host=foo",), 49 | ("postgresql://admin:pass@foo:5432/db",), 50 | ], 51 | ) 52 | def test_connect_raises_connection_error(conn_str: tuple[str]) -> None: 53 | with pytest.raises(HarlequinConnectionError): 54 | _ = HarlequinPostgresAdapter(conn_str=conn_str, connect_timeout=0.1).connect() 55 | 56 | 57 | @pytest.mark.parametrize( 58 | "conn_str,options,expected", 59 | [ 60 | (("",), {}, "localhost:5432/postgres"), 61 | (("host=foo",), {}, "foo:5432/postgres"), 62 | (("postgresql://foo",), {}, "foo:5432/postgres"), 63 | (("postgresql://foo",), {"port": 5431}, "foo:5431/postgres"), 64 | (("postgresql://foo/mydb",), {"port": 5431}, "foo:5431/mydb"), 65 | (("postgresql://admin:pass@foo/mydb",), {"port": 5431}, "foo:5431/mydb"), 66 | (("postgresql://admin:pass@foo:5431/mydb",), {}, "foo:5431/mydb"), 67 | ], 68 | ) 69 | def test_connection_id( 70 | conn_str: tuple[str], options: dict[str, int | float | str | None], expected: str 71 | ) -> None: 72 | adapter = HarlequinPostgresAdapter( 73 | conn_str=conn_str, 74 | **options, # type: ignore[arg-type] 75 | ) 76 | assert adapter.connection_id == expected 77 | 78 | 79 | def test_get_catalog(connection: HarlequinPostgresConnection) -> None: 80 | catalog = connection.get_catalog() 81 | assert isinstance(catalog, Catalog) 82 | assert catalog.items 83 | assert isinstance(catalog.items[0], CatalogItem) 84 | 85 | 86 | def test_get_completions(connection: HarlequinPostgresConnection) -> None: 87 | completions = connection.get_completions() 88 | test_labels = ["atomic", "greatest", "point_right", "autovacuum"] 89 | filtered = list(filter(lambda x: x.label in test_labels, completions)) 90 | assert len(filtered) == 4 91 | value_filtered = list(filter(lambda x: x.value in test_labels, completions)) 92 | assert len(value_filtered) == 4 93 | 94 | 95 | def test_execute_ddl(connection: HarlequinPostgresConnection) -> None: 96 | cur = connection.execute("create table foo (a int)") 97 | assert cur is None 98 | 99 | 100 | def test_execute_select(connection: HarlequinPostgresConnection) -> None: 101 | cur = connection.execute("select 1 as a") 102 | assert isinstance(cur, HarlequinCursor) 103 | assert cur.columns() == [("a", "#")] 104 | data = cur.fetchall() 105 | backend = create_backend(data) 106 | assert backend.column_count == 1 107 | assert backend.row_count == 1 108 | 109 | 110 | def test_execute_select_dupe_cols(connection: HarlequinPostgresConnection) -> None: 111 | cur = connection.execute("select 1 as a, 2 as a, 3 as a") 112 | assert isinstance(cur, HarlequinCursor) 113 | assert len(cur.columns()) == 3 114 | data = cur.fetchall() 115 | backend = create_backend(data) 116 | assert backend.column_count == 3 117 | assert backend.row_count == 1 118 | 119 | 120 | def test_set_limit(connection: HarlequinPostgresConnection) -> None: 121 | cur = connection.execute("select 1 as a union all select 2 union all select 3") 122 | assert isinstance(cur, HarlequinCursor) 123 | cur = cur.set_limit(2) 124 | assert isinstance(cur, HarlequinCursor) 125 | data = cur.fetchall() 126 | backend = create_backend(data) 127 | assert backend.column_count == 1 128 | assert backend.row_count == 2 129 | 130 | 131 | def test_execute_raises_query_error(connection: HarlequinPostgresConnection) -> None: 132 | with pytest.raises(HarlequinQueryError): 133 | _ = connection.execute("sel;") 134 | 135 | 136 | def test_inf_timestamps(connection: HarlequinPostgresConnection) -> None: 137 | cur = connection.execute( 138 | """select 139 | 'infinity'::date, 140 | 'infinity'::timestamp, 141 | 'infinity'::timestamptz, 142 | '-infinity'::date, 143 | '-infinity'::timestamp, 144 | '-infinity'::timestamptz 145 | """ 146 | ) 147 | assert cur is not None 148 | data = cur.fetchall() 149 | assert data == [ 150 | ( 151 | date.max, 152 | datetime.max, 153 | datetime.max, 154 | date.min, 155 | datetime.min, 156 | datetime.min, 157 | ) 158 | ] 159 | -------------------------------------------------------------------------------- /tests/test_catalog.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from harlequin.catalog import InteractiveCatalogItem 3 | from harlequin_postgres.adapter import HarlequinPostgresConnection 4 | from harlequin_postgres.catalog import ( 5 | ColumnCatalogItem, 6 | DatabaseCatalogItem, 7 | RelationCatalogItem, 8 | SchemaCatalogItem, 9 | TableCatalogItem, 10 | ViewCatalogItem, 11 | ) 12 | 13 | 14 | @pytest.fixture 15 | def connection_with_objects( 16 | connection: HarlequinPostgresConnection, 17 | ) -> HarlequinPostgresConnection: 18 | connection.execute("create schema one") 19 | connection.execute("create table one.foo as select 1 as a, '2' as b") 20 | connection.execute("create table one.bar as select 1 as a, '2' as b") 21 | connection.execute("create table one.baz as select 1 as a, '2' as b") 22 | connection.execute("create schema two") 23 | connection.execute("create view two.qux as select * from one.foo") 24 | connection.execute("create schema three") 25 | # the original connection fixture will clean this up. 26 | return connection 27 | 28 | 29 | def test_catalog(connection_with_objects: HarlequinPostgresConnection) -> None: 30 | conn = connection_with_objects 31 | 32 | catalog = conn.get_catalog() 33 | 34 | # at least two databases, postgres and test 35 | assert len(catalog.items) >= 2 36 | 37 | [test_db_item] = filter(lambda item: item.label == "test", catalog.items) 38 | assert isinstance(test_db_item, InteractiveCatalogItem) 39 | assert isinstance(test_db_item, DatabaseCatalogItem) 40 | assert not test_db_item.children 41 | assert not test_db_item.loaded 42 | 43 | schema_items = test_db_item.fetch_children() 44 | assert all(isinstance(item, SchemaCatalogItem) for item in schema_items) 45 | 46 | [schema_one_item] = filter(lambda item: item.label == "one", schema_items) 47 | assert isinstance(schema_one_item, SchemaCatalogItem) 48 | assert not schema_one_item.children 49 | assert not schema_one_item.loaded 50 | 51 | table_items = schema_one_item.fetch_children() 52 | assert all(isinstance(item, RelationCatalogItem) for item in table_items) 53 | 54 | [foo_item] = filter(lambda item: item.label == "foo", table_items) 55 | assert isinstance(foo_item, TableCatalogItem) 56 | assert not foo_item.children 57 | assert not foo_item.loaded 58 | 59 | foo_column_items = foo_item.fetch_children() 60 | assert all(isinstance(item, ColumnCatalogItem) for item in foo_column_items) 61 | 62 | [schema_two_item] = filter(lambda item: item.label == "two", schema_items) 63 | assert isinstance(schema_two_item, SchemaCatalogItem) 64 | assert not schema_two_item.children 65 | assert not schema_two_item.loaded 66 | 67 | view_items = schema_two_item.fetch_children() 68 | assert all(isinstance(item, ViewCatalogItem) for item in view_items) 69 | 70 | [qux_item] = filter(lambda item: item.label == "qux", view_items) 71 | assert isinstance(qux_item, ViewCatalogItem) 72 | assert not qux_item.children 73 | assert not qux_item.loaded 74 | 75 | qux_column_items = qux_item.fetch_children() 76 | assert all(isinstance(item, ColumnCatalogItem) for item in qux_column_items) 77 | 78 | assert [item.label for item in foo_column_items] == [ 79 | item.label for item in qux_column_items 80 | ] 81 | 82 | # ensure calling fetch_children on cols doesn't raise 83 | children_items = foo_column_items[0].fetch_children() 84 | assert not children_items 85 | 86 | [schema_three_item] = filter(lambda item: item.label == "three", schema_items) 87 | assert isinstance(schema_two_item, SchemaCatalogItem) 88 | assert not schema_two_item.children 89 | assert not schema_two_item.loaded 90 | 91 | three_children = schema_three_item.fetch_children() 92 | assert not three_children 93 | --------------------------------------------------------------------------------