├── .gitignore ├── LICENSE ├── MANIFEST.in ├── README.md ├── hatch_build.py ├── images ├── console.png └── notebook.png ├── postgres_kernel ├── __init__.py ├── __main__.py ├── _version.py └── kernel.py └── pyproject.toml /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | postgres_kernel.egg-info 3 | *.pyc 4 | build/ 5 | dist/ 6 | MANIFEST 7 | *.ipynb_checkpoints* 8 | data_kernelspec 9 | 10 | \.vscode/ 11 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2016 Brian Schiller 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 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include README.rst 2 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # A simple Jupyter kernel for PostgreSQL 2 | 3 | Install with `pip install postgres_kernel` 4 | 5 | To use, run one of: 6 | 7 | ```bash 8 | jupyter notebook 9 | # In the notebook interface, select PostgreSQL from the 'New' menu 10 | jupyter qtconsole --kernel postgres 11 | jupyter console --kernel postgres 12 | ``` 13 | 14 | ## How to use: 15 | 16 | There are a couple of specially formatted comments for controlling the connection string and autocommit mode. 17 | 18 | ```sql 19 | -- connection: postgres://brian:password@localhost:5432/dbname 20 | -- autocommit: true 21 | -- (or false) 22 | ``` 23 | 24 | For details of how this works, see Jupyter's docs on [wrapper kernels](http://jupyter-client.readthedocs.io/en/latest/wrapperkernels.html). 25 | This is heavily based on [takluyver/bash_kernel](https://github.com/takluyver/bash_kernel). Just look at our git log :) 26 | 27 | ![](images/console.png) 28 | 29 | ![](images/notebook.png) 30 | 31 | 32 | Related 33 | ------- 34 | 35 | - Catherine Devlin has an ipython magic that seems very full featured: [catherinedevlin/ipython-sql](https://github.com/catherinedevlin/ipython-sql) 36 | 37 | - As noted, this is based on [takluyver/bash_kernel](https://github.com/takluyver/bash_kernel) 38 | -------------------------------------------------------------------------------- /hatch_build.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | from pathlib import Path 3 | from ipykernel.kernelspec import make_ipkernel_cmd, write_kernel_spec 4 | from hatchling.builders.hooks.plugin.interface import BuildHookInterface 5 | 6 | 7 | class CustomHook(BuildHookInterface): 8 | 9 | def initialize(self, version, build_data): 10 | dest = Path(__file__).parent.resolve() / "data_kernelspec" 11 | overrides = { 12 | "argv": make_ipkernel_cmd( 13 | executable="python", mod='postgres_kernel' 14 | ), 15 | "display_name": "PosgreSQL", 16 | "language": "sql", 17 | "codemirror_mode": "sql" 18 | } 19 | 20 | if dest.exists(): 21 | shutil.rmtree(dest) 22 | 23 | write_kernel_spec(dest, overrides=overrides) 24 | -------------------------------------------------------------------------------- /images/console.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bgschiller/postgres_kernel/1bb3905b2a5ab3dbb5ec2bad341bbf8320579e7a/images/console.png -------------------------------------------------------------------------------- /images/notebook.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bgschiller/postgres_kernel/1bb3905b2a5ab3dbb5ec2bad341bbf8320579e7a/images/notebook.png -------------------------------------------------------------------------------- /postgres_kernel/__init__.py: -------------------------------------------------------------------------------- 1 | from ._version import __version__ 2 | -------------------------------------------------------------------------------- /postgres_kernel/__main__.py: -------------------------------------------------------------------------------- 1 | from ipykernel.kernelapp import IPKernelApp 2 | from .kernel import PostgresKernel 3 | IPKernelApp.launch_instance(kernel_class=PostgresKernel) 4 | -------------------------------------------------------------------------------- /postgres_kernel/_version.py: -------------------------------------------------------------------------------- 1 | __version__ = '0.2.3' 2 | -------------------------------------------------------------------------------- /postgres_kernel/kernel.py: -------------------------------------------------------------------------------- 1 | from ipykernel.kernelbase import Kernel 2 | import psycopg2 3 | from psycopg2 import Error, ProgrammingError, OperationalError 4 | from psycopg2.extensions import ( 5 | QueryCanceledError, POLL_OK, POLL_READ, POLL_WRITE, STATUS_BEGIN, 6 | ) 7 | 8 | import re 9 | import os 10 | from select import select 11 | 12 | from ._version import __version__ 13 | from tabulate import tabulate 14 | version_pat = re.compile(r'^PostgreSQL (\d+(\.\d+)+)') 15 | 16 | 17 | def log(val): 18 | return # comment out line for debug 19 | with open('kernel.log', 'a') as f: 20 | f.write(str(val) + '\n') 21 | return val 22 | 23 | 24 | def wait_select_inter(conn): 25 | while 1: 26 | try: 27 | state = conn.poll() 28 | if state == POLL_OK: 29 | break 30 | elif state == POLL_READ: 31 | select([conn.fileno()], [], []) 32 | elif state == POLL_WRITE: 33 | select([], [conn.fileno()], []) 34 | else: 35 | raise conn.OperationalError( 36 | "bad state from poll: %s" % state) 37 | except KeyboardInterrupt: 38 | conn.cancel() 39 | # the loop will be broken by a server error 40 | continue 41 | 42 | 43 | class PostgresKernel(Kernel): 44 | implementation = 'postgres_kernel' 45 | implementation_version = __version__ 46 | 47 | language_info = {'name': 'PostgreSQL', 48 | 'codemirror_mode': 'sql', 49 | 'mimetype': 'text/x-postgresql', 50 | 'file_extension': '.sql'} 51 | 52 | def __init__(self, **kwargs): 53 | Kernel.__init__(self, **kwargs) 54 | log('inside init') 55 | # Catch KeyboardInterrupt, cancel query, raise QueryCancelledError 56 | psycopg2.extensions.set_wait_callback(wait_select_inter) 57 | self._conn_string = os.getenv('DATABASE_URL', '') 58 | self._autocommit = True 59 | self._conn = None 60 | self._start_connection() 61 | 62 | @property 63 | def language_version(self): 64 | m = version_pat.search(self.banner) 65 | return m.group(1) 66 | 67 | _banner = None 68 | 69 | @property 70 | def banner(self): 71 | if self._banner is None: 72 | if self._conn is None: 73 | return 'not yet connected to a database' 74 | self._banner = self.fetchone('SELECT VERSION();')[0] 75 | return self._banner 76 | 77 | def _start_connection(self): 78 | log('starting connection') 79 | try: 80 | self._conn = psycopg2.connect(self._conn_string) 81 | self._conn.autocommit = self._autocommit 82 | except OperationalError: 83 | log('failed to connect to {}'.format(self._conn_string)) 84 | message = '''Failed to connect to a database at {}'''.format(self._conn_string) 85 | self.send_response(self.iopub_socket, 'stream', 86 | {'name': 'stderr', 'text': message}) 87 | 88 | def fetchone(self, query): 89 | log('fetching one from: \n' + query) 90 | with self._conn.cursor() as c: 91 | c.execute(query) 92 | one = c.fetchone() 93 | log(one) 94 | return one 95 | 96 | def fetchall(self, query): 97 | log('fetching all from: \n' + query) 98 | with self._conn.cursor() as c: 99 | c.execute(query) 100 | desc = c.description 101 | if c.description: 102 | keys = [col[0] for col in c.description] 103 | return keys, c.fetchall() 104 | return None, None 105 | 106 | CONN_STRING_COMMENT = re.compile(r'--\s*connection:\s*(.*)\s*$') 107 | AUTOCOMMIT_SWITCH_COMMENT = re.compile(r'--\s*autocommit:\s*(\w+)\s*$') 108 | 109 | def change_connection(self, conn_string): 110 | self._conn_string = conn_string 111 | self._start_connection() 112 | 113 | def switch_autocommit(self, switch_to): 114 | self._autocommit = switch_to 115 | committed = False 116 | if self._conn: 117 | if self._conn.get_transaction_status() == STATUS_BEGIN: 118 | committed = True 119 | self._conn.commit() 120 | self._conn.autocommit = switch_to 121 | else: 122 | self._start_connection() 123 | return committed 124 | 125 | def change_autocommit_mode(self, switch): 126 | """ 127 | Strip and make a string case insensitive and ensure it is either 'true' or 'false'. 128 | 129 | If neither, prompt user for either value. 130 | When 'true', return True, and when 'false' return False. 131 | """ 132 | parsed_switch = switch.strip().lower() 133 | if not parsed_switch in ['true', 'false']: 134 | self.send_response( 135 | self.iopub_socket, 'stream', { 136 | 'name': 'stderr', 137 | 'text': 'autocommit must be true or false.\n\n' 138 | } 139 | ) 140 | 141 | switch_bool = (parsed_switch == 'true') 142 | committed = self.switch_autocommit(switch_bool) 143 | message = ( 144 | 'committed current transaction & ' if committed else '' + 145 | 'switched autocommit mode to ' + 146 | str(self._autocommit) 147 | ) 148 | self.send_response( 149 | self.iopub_socket, 'stream', { 150 | 'name': 'stderr', 151 | 'text': message, 152 | } 153 | ) 154 | 155 | def do_execute(self, code, silent, store_history=True, 156 | user_expressions=None, allow_stdin=False): 157 | print(code) 158 | 159 | connection_string = self.CONN_STRING_COMMENT.findall(code) 160 | autocommit_switch = self.AUTOCOMMIT_SWITCH_COMMENT.findall(code) 161 | if autocommit_switch: 162 | self.change_autocommit_mode(autocommit_switch[0]) 163 | if connection_string: 164 | self.change_connection(connection_string[0]) 165 | 166 | code = self.AUTOCOMMIT_SWITCH_COMMENT.sub('', self.CONN_STRING_COMMENT.sub('', code)) 167 | if not code.strip(): 168 | return {'status': 'ok', 'execution_count': self.execution_count, 169 | 'payload': [], 'user_expressions': {}} 170 | 171 | if self._conn is None: 172 | self.send_response( 173 | self.iopub_socket, 'stream', { 174 | 'name': 'stderr', 175 | 'text': '''\ 176 | Error: Unable to connect to a database at "{}". 177 | Perhaps you need to set a connection string with 178 | -- connection: '''.format(self._conn_string) 179 | }) 180 | return {'status': 'error', 'execution_count': self.execution_count, 181 | 'ename': 'MissingConnection'} 182 | try: 183 | header, rows = self.fetchall(code) 184 | except QueryCanceledError: 185 | self._conn.rollback() 186 | return {'status': 'abort', 'execution_count': self.execution_count} 187 | except Error as e: 188 | self.send_response(self.iopub_socket, 'stream', 189 | {'name': 'stderr', 'text': str(e)}) 190 | self._conn.rollback() 191 | return {'status': 'error', 'execution_count': self.execution_count, 192 | 'ename': 'ProgrammingError', 'evalue': str(e), 193 | 'traceback': []} 194 | else: 195 | if rows is not None: 196 | self.send_response( 197 | self.iopub_socket, 'stream', { 198 | 'name': 'stdout', 199 | 'text': str(len(rows)) + " row(s) returned.\n" 200 | }) 201 | 202 | for notice in self._conn.notices: 203 | self.send_response( 204 | self.iopub_socket, 'stream', { 205 | 'name': 'stdout', 206 | 'text': str(notice) 207 | }) 208 | self._conn.notices = [] 209 | 210 | if header is not None and len(rows) > 0: 211 | self.send_response(self.iopub_socket, 'display_data', display_data(header, rows)) 212 | 213 | return {'status': 'ok', 'execution_count': self.execution_count, 214 | 'payload': [], 'user_expressions': {}} 215 | 216 | 217 | def display_data(header, rows): 218 | d = { 219 | 'data': { 220 | 'text/latex': tabulate(rows, header, tablefmt='latex_booktabs'), 221 | 'text/plain': tabulate(rows, header, tablefmt='simple'), 222 | 'text/html': tabulate(rows, header, tablefmt='html'), 223 | }, 224 | 'metadata': {} 225 | } 226 | return d 227 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling", "ipykernel"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "postgres_kernel" 7 | dynamic = ["version"] 8 | authors = [{name = "Brian Schiller", email = "bgschiller@gmail.com"}] 9 | dependencies = [ 10 | 'psycopg2>=2.6', 11 | 'tabulate>=0.7.5', 12 | 'jupyter-client', 13 | 'ipykernel' 14 | ] 15 | classifiers = [ 16 | 'Framework :: IPython', 17 | 'License :: OSI Approved :: BSD License', 18 | 'Programming Language :: Python :: 3', 19 | 'Topic :: System :: Shells', 20 | ] 21 | readme = "README.md" 22 | 23 | [project.urls] 24 | Source = "https://github.com/bgschiller/postgres_kernel" 25 | 26 | [tool.hatch.version] 27 | path = "postgres_kernel/_version.py" 28 | 29 | # Used to call hatch_build.py 30 | [tool.hatch.build.hooks.custom] 31 | 32 | [tool.hatch.build.targets.wheel.shared-data] 33 | "data_kernelspec" = "share/jupyter/kernels/postgres_kernel" --------------------------------------------------------------------------------