├── .gitignore ├── README.md ├── pyproject.toml └── sqlalchemy_fdw ├── __init__.py ├── dialect.py └── util.py /.gitignore: -------------------------------------------------------------------------------- 1 | dist/ 2 | __pycache__/ 3 | *.egg-info 4 | .venv/ 5 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | sqlalchemy_fdw is a postgresql dialect for sqlalchemy adding support for foreign 2 | data wrappers. 3 | 4 | Installation 5 | ------------ 6 | ```bash 7 | 8 | pip install sqlalchemy-fdw 9 | ``` 10 | 11 | Usage 12 | ----- 13 | 14 | ```python 15 | from sqlalchemy import create_engine, Table, Column, MetaData 16 | from sqlalchemy import Integer, Unicode 17 | from sqlalchemy_fdw import ForeignTable, ForeignDataWrapper 18 | 19 | 20 | engine = create_engine('pgfdw://user:password@host:port/dbname') 21 | metadata = MetaData() 22 | metadata.bind = engine 23 | 24 | fdw = ForeignDataWrapper("myfdwserver", "myfdwextension", metadata=metadata, 25 | options={'option1': 'test'}) 26 | fdw.create() 27 | 28 | table = ForeignTable("myforeigntable", metadata, 29 | Column('col1', Integer), 30 | Column('col2', Unicode), 31 | pgfdw_server='myfdwserver', 32 | pgfdw_options={ 33 | 'tableoption': 'optionvalue' 34 | } 35 | ) 36 | table.create(checkfirst=True) 37 | table.drop() 38 | fdw.drop(cascade=True) 39 | ``` 40 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ['flit_core >=2,<4'] 3 | build-backend = 'flit_core.buildapi' 4 | 5 | [project] 6 | name = 'sqlalchemy_fdw' 7 | version = '0.4.1' 8 | description = 'PostgreSQL dialect supporting foreign data wrappers' 9 | authors = [{name = 'Kozea'}] 10 | license = {text = 'BSD'} 11 | requires-python = '>=3.6' 12 | dependencies = [ 13 | 'sqlalchemy>=1.3', 14 | 'psycopg2-binary', 15 | ] 16 | 17 | [project.entry-points."sqlalchemy.dialects"] 18 | pgfdw = 'sqlalchemy_fdw.dialect:dialect' 19 | -------------------------------------------------------------------------------- /sqlalchemy_fdw/__init__.py: -------------------------------------------------------------------------------- 1 | """Contains Schema element and compilers for foreign table and fdw. 2 | """ 3 | 4 | from sqlalchemy.ext.compiler import compiles 5 | from sqlalchemy.schema import DDLElement, Table 6 | from sqlalchemy.sql.base import _bind_or_error 7 | from sqlalchemy import sql 8 | from sqlalchemy import types 9 | from .util import sql_options 10 | 11 | 12 | class ForeignTable(Table): 13 | """Defines a Foreign Table 14 | 15 | A Foreign Table is a postgresql table located on a remote server. 16 | To create remote servers, look at :class:`ForeignDataWrapper`. 17 | 18 | This functionality has been tagged stable in postgresql 9.1 19 | 20 | Assuming you already created a server 'myserver', a foreign table can be 21 | defined on it like this:: 22 | 23 | mytable = ForeignTable("mytable", metadata, 24 | Column('id', Integer), 25 | Column('name', Unicode), 26 | pgfdw_server='myserver) 27 | 28 | You can then use it like any table, except: 29 | - only select statements are supported 30 | - constraints are not supported 31 | 32 | Constructor arguments are the same as :class:`Table`, plus: 33 | 34 | :param: pgfdw_server: the name of the server this table belongs to. 35 | 36 | :param: pgfdw_options: a dictionary containing the table options. 37 | 38 | These options are passed directly to the foreign table as an OPTIONS 39 | clause. 40 | 41 | e.g:: 42 | 43 | mytable = ForeignTable("mytable", metadata, 44 | Column('id', Integer), 45 | Column('name', Unicode), 46 | pgfdw_server='myserver', 47 | pgfdw_options={'option1': 'test'}) 48 | 49 | Results in the following sql:: 50 | 51 | CREATE FOREIGN TABLE mytable ( 52 | id integer, 53 | name character varying 54 | ) server myserver options ( 55 | "option1" 'test' 56 | ); 57 | 58 | 59 | For more information on the available foreign data wrappers, 60 | see `http://pgxn.org/tag/foreign%20data%20wrapper/`. 61 | 62 | """ 63 | 64 | def __new__(cls, *args, **kwargs): 65 | if not args: 66 | # python3k pickle seems to call this 67 | return object.__new__(cls) 68 | 69 | table = super(ForeignTable, cls).__new__(cls, *args, **kwargs) 70 | metadata = args[1] 71 | table.pgfdw_server = kwargs.pop('pgfdw_server', None) 72 | table.pgfdw_options = kwargs.pop('pgfdw_options', None) or {} 73 | if 'FOREIGN' not in table._prefixes: 74 | table._prefixes.append('FOREIGN') 75 | 76 | if not hasattr(metadata, '_foreign_tables'): 77 | metadata._foreign_tables = {} 78 | 79 | metadata._foreign_tables[table.key] = table 80 | 81 | autoload = kwargs.get('autoload', False) 82 | autoload_with = kwargs.get('autoload_with', None) 83 | if autoload: 84 | if autoload_with: 85 | autoload_with.run_callable( 86 | autoload_with.dialect.get_foreign_table_options, 87 | table) 88 | else: 89 | bind = _bind_or_error( 90 | table.metadata, 91 | msg="No engine is bound to this ForeignTable's MetaData. " 92 | "Pass an engine to the Table via " 93 | "autoload_with=, " 94 | "or associate the MetaData with an engine via " 95 | "metadata.bind=") 96 | bind.run_callable( 97 | bind.dialect.get_foreign_table_options, table) 98 | return table 99 | 100 | 101 | class ForeignDataWrapper(DDLElement): 102 | """Defines a foreign data wrapper server 103 | 104 | A foreign data wrapper server must be defined to access a foreign data 105 | wrapper installed as an extension. 106 | 107 | Basic usage:: 108 | 109 | myfdw = ForeignDataWrapper('myfdw', 'mysql_fdw', metadata) 110 | myfdw.create(checkfirst=True) 111 | # Create some :class:`ForeignTable` 112 | myfdw.drop() 113 | 114 | Constructor accepts the following arguments: 115 | 116 | :param: name: the server name 117 | :param: extension_name: the foreign data wrapper extension to be used 118 | :param: metadata: (optional) the :class:`MetaData` object to bind with 119 | :param: bind: (optional) the :class:`Engine` object to bind with 120 | 121 | """ 122 | 123 | def __init__(self, name, extension_name, metadata=None, bind=None, 124 | options=None): 125 | self.name = name 126 | self.options = options or {} 127 | self.extension_name = extension_name 128 | self.metadata = metadata 129 | self._bind = bind 130 | 131 | @property 132 | def bind(self): 133 | """Returns the current bind""" 134 | return self._bind or self.metadata.bind 135 | 136 | def check_existence(self, bind=None): 137 | """Checks if a server with the same name already exists. 138 | 139 | :param: bind: (optional) if not bind is supplied, the current binding 140 | (from the metatadata) will be used. 141 | 142 | """ 143 | if bind is None: 144 | bind = _bind_or_error(self) 145 | bindparams = [ 146 | sql.bindparam('name', str(self.name), type_=types.Unicode)] 147 | cursor = bind.execute( 148 | sql.text( 149 | "select srvname from pg_foreign_server where srvname = :name" 150 | ).bindparams(*bindparams) 151 | ) 152 | return bool(cursor.first()) 153 | 154 | def create(self, bind=None, checkfirst=False): 155 | """Create the server. 156 | 157 | :param: bind: (optional) The bind to use instead of the instance one 158 | :param: checkfirst: Check if the server exists before creating it. 159 | 160 | """ 161 | 162 | if bind is None: 163 | bind = _bind_or_error(self) 164 | if not checkfirst or not self.check_existence(bind): 165 | CreateForeignDataWrapper( 166 | self.name, self.extension_name, bind=bind, options=self.options 167 | ).execute() 168 | 169 | def drop(self, bind=None, checkfirst=False, cascade=False): 170 | """Drop the server 171 | 172 | :param: bind: (optional) The bind to use instead of the instance one 173 | :param: checkfirst: Check if the server exists before dropping it. 174 | :param: cascade: appends the CASCADE keyword to the drop statement. 175 | 176 | """ 177 | if bind is None: 178 | bind = _bind_or_error(self) 179 | if not checkfirst or self.check_existence(bind): 180 | DropForeignDataWrapper( 181 | self.name, self.extension_name, bind=bind, cascade=cascade 182 | ).execute() 183 | 184 | 185 | class CreateForeignDataWrapper(ForeignDataWrapper): 186 | """The concrete create statement""" 187 | pass 188 | 189 | 190 | class DropForeignDataWrapper(ForeignDataWrapper): 191 | """The concrete drop statement""" 192 | 193 | def __init__(self, *args, **kwargs): 194 | self.cascade = kwargs.pop('cascade', False) 195 | super(DropForeignDataWrapper, self).__init__(*args, **kwargs) 196 | 197 | 198 | @compiles(CreateForeignDataWrapper) 199 | def visit_create_fdw(create, compiler, **kw): 200 | """Compiler for the create server statement""" 201 | preparer = compiler.dialect.identifier_preparer 202 | statement = ( 203 | "CREATE server %s foreign data wrapper %s " % ( 204 | preparer.quote_identifier(create.name), 205 | preparer.quote_identifier(create.extension_name))) 206 | statement += sql_options(create.options, preparer) 207 | return statement 208 | 209 | 210 | @compiles(DropForeignDataWrapper) 211 | def visit_drop_fdw(drop, compiler, **kw): 212 | """Compiler for drop server statement""" 213 | preparer = compiler.dialect.identifier_preparer 214 | statement = "DROP server %s " % (preparer.quote_identifier(drop.name)) 215 | if drop.cascade: 216 | statement += " CASCADE" 217 | return statement 218 | -------------------------------------------------------------------------------- /sqlalchemy_fdw/dialect.py: -------------------------------------------------------------------------------- 1 | """A custom dialect for handling foreign tables on postgresql""" 2 | 3 | from sqlalchemy.dialects.postgresql import ARRAY 4 | from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2 5 | from sqlalchemy.dialects.postgresql.base import PGDDLCompiler 6 | from sqlalchemy.engine import reflection 7 | from sqlalchemy.schema import Table, ForeignKeyConstraint 8 | from sqlalchemy import sql, types as sqltypes, exc 9 | from .util import sql_options 10 | 11 | 12 | def is_foreign(t): 13 | return t.key in getattr(t.metadata, '_foreign_tables', {}) 14 | 15 | 16 | class PGDDLCompilerFdw(PGDDLCompiler): 17 | """A DDL compiler for the pgfdw dialect, for managing foreign tables""" 18 | 19 | def post_create_table(self, table): 20 | if is_foreign(table): 21 | preparer = self.dialect.identifier_preparer 22 | post = ' server %s ' % table.pgfdw_server 23 | post += sql_options(table.pgfdw_options, preparer) 24 | return post 25 | else: 26 | return super(PGDDLCompilerFdw, self).post_create_table(table) 27 | 28 | def visit_drop_table(self, drop): 29 | prefix = "" 30 | if is_foreign(drop.element): 31 | prefix = "FOREIGN" 32 | return "DROP %s TABLE %s" % ( 33 | prefix, self.preparer.format_table(drop.element)) 34 | 35 | def create_table_constraints(self, table, 36 | _include_foreign_key_constraints=None): 37 | # No constraint in foreign tables 38 | if is_foreign(table): 39 | return '' 40 | else: 41 | constraints = [] 42 | if table.primary_key: 43 | constraints.append(table.primary_key) 44 | 45 | constraints.extend([c for c in table._sorted_constraints 46 | if c is not table.primary_key]) 47 | 48 | def foreign_foreign_key(constraint): 49 | """Return whether this is a foreign key 50 | referencing a foreign table""" 51 | return isinstance( 52 | constraint, ForeignKeyConstraint 53 | ) and is_foreign(constraint.referred_table) 54 | 55 | return ", \n\t".join(p for p in ( 56 | self.process(constraint) 57 | for constraint in constraints 58 | if ( 59 | constraint._create_rule is None or 60 | constraint._create_rule(self)) 61 | and ( 62 | not self.dialect.supports_alter or 63 | not getattr(constraint, 'use_alter', False)) 64 | and not foreign_foreign_key(constraint) 65 | ) if p is not None) 66 | 67 | 68 | class PGDialectFdw(PGDialect_psycopg2): 69 | """An sqldialect based on psyopg2 for managing foreign tables 70 | 71 | To use it, simply use pgfdw in the connection string:: 72 | 73 | create_engine('pgfdw://user:password@localhost:5432/dbname') 74 | 75 | """ 76 | supports_statement_cache = True 77 | ddl_compiler = PGDDLCompilerFdw 78 | 79 | construct_arguments = [ 80 | (Table, { 81 | "server": None, 82 | "options": None 83 | }) 84 | ] 85 | 86 | @reflection.cache 87 | def get_primary_keys(self, connection, table_name, schema=None, **kw): 88 | if schema is not None: 89 | current_schema = schema 90 | else: 91 | current_schema = self.default_schema_name 92 | PK_SQL = """ 93 | SELECT cu.column_name 94 | FROM information_schema.table_constraints tc 95 | INNER JOIN information_schema.key_column_usage cu 96 | on cu.constraint_name = tc.constraint_name and 97 | cu.table_name = tc.table_name and 98 | cu.table_schema = tc.table_schema 99 | WHERE cu.table_name = :table_name and 100 | constraint_type = 'PRIMARY KEY' 101 | and cu.table_schema = :schema; 102 | """ 103 | t = sql.text(PK_SQL).columns(attname=sqltypes.Unicode) 104 | c = connection.execute(t, table_name=table_name, schema=current_schema) 105 | primary_keys = [r[0] for r in c.fetchall()] 106 | return primary_keys 107 | 108 | @reflection.cache 109 | def get_table_names(self, connection, schema=None, **kw): 110 | if schema is not None: 111 | current_schema = schema 112 | else: 113 | current_schema = self.default_schema_name 114 | 115 | result = connection.execute( 116 | sql.text( 117 | "SELECT relname FROM pg_class c " 118 | "WHERE relkind in ('r', 'f') " 119 | f"AND '{current_schema}' = (select nspname from pg_namespace n " 120 | "where n.oid = c.relnamespace) " 121 | ).columns(relname=sqltypes.Unicode) 122 | ) 123 | return [row[0] for row in result] 124 | 125 | @reflection.cache 126 | def get_table_oid(self, connection, table_name, schema=None, **kw): 127 | """Fetch the oid for schema.table_name. 128 | 129 | Several reflection methods require the table oid. The idea for using 130 | this method is that it can be fetched one time and cached for 131 | subsequent calls. 132 | 133 | """ 134 | table_oid = None 135 | if schema is not None: 136 | schema_where_clause = "n.nspname = :schema" 137 | else: 138 | schema_where_clause = "pg_catalog.pg_table_is_visible(c.oid)" 139 | query = """ 140 | SELECT c.oid 141 | FROM pg_catalog.pg_class c 142 | LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace 143 | WHERE (%s) 144 | AND c.relname = :table_name AND c.relkind in ('r', 'v', 'f') 145 | """ % schema_where_clause 146 | # Since we're binding to unicode, table_name and schema_name must be 147 | # unicode. 148 | table_name = str(table_name) 149 | bindparams = [sql.bindparam('table_name', type_=sqltypes.Unicode)] 150 | if schema is not None: 151 | schema = str(schema) 152 | bindparams.append(sql.bindparam('schema', type_=sqltypes.Unicode)) 153 | s = sql.text(query).bindparams(*bindparams).columns(oid=sqltypes.Integer) 154 | c = connection.execute(s, table_name=table_name, schema=schema) 155 | table_oid = c.scalar() 156 | if table_oid is None: 157 | raise exc.NoSuchTableError(table_name) 158 | return table_oid 159 | 160 | @reflection.cache 161 | def get_foreign_table_options(self, connection, pgfdw_table): 162 | oid = self.get_table_oid(connection, pgfdw_table.name, 163 | pgfdw_table.schema) 164 | query = """ 165 | SELECT ftoptions, srvname 166 | FROM pg_foreign_table t inner join pg_foreign_server s 167 | ON t.ftserver = s.oid 168 | WHERE t.ftrelid = :oid 169 | """ 170 | s = sql.text(query).bindparams( 171 | sql.bindparam('oid', type_=sqltypes.Integer) 172 | ).columns(ftoptions=ARRAY(sqltypes.Unicode), srvname=sqltypes.Unicode) 173 | c = connection.execute(s, oid=oid) 174 | options, srv_name = c.fetchone() 175 | pgfdw_table.pgfdw_server = srv_name 176 | pgfdw_table.pgfdw_options = dict([ 177 | option.split('=', 1) for option in options 178 | ]) if options is not None else {} 179 | 180 | dialect = PGDialectFdw 181 | -------------------------------------------------------------------------------- /sqlalchemy_fdw/util.py: -------------------------------------------------------------------------------- 1 | """Utilities used by both the dialect and the schema objects""" 2 | 3 | 4 | def sql_options(options, preparer): 5 | """Format an options clause, if any""" 6 | if options: 7 | return ' options (%s)' % ','.join([ 8 | "%s '%s'" % (preparer.quote_identifier(key), value) 9 | for key, value in list(options.items()) 10 | ]) 11 | return '' 12 | --------------------------------------------------------------------------------