Merge pull request #30 from xinbo/master

sqlalchemy: Support references to tables in other schemas

cockroachdb/sqlalchemy/dialect.py

@@ -8,6 +8,7 @@
 import sqlalchemy.types as sqltypes
 
 from .ddl_compiler import DDLCompiler
+from .stmt_compiler import CockroachCompiler
 
 # Map type names (as returned by SHOW COLUMNS) to sqlalchemy type
 # objects.
@@ -58,6 +59,7 @@ class CockroachDBDialect(PGDialect_psycopg2):
     name = 'cockroachdb'
     supports_sequences = False
     ddl_compiler = DDLCompiler
+    statement_compiler = CockroachCompiler
 
     def __init__(self, *args, **kwargs):
         super(CockroachDBDialect, self).__init__(*args,
@@ -97,7 +99,7 @@ def has_table(self, conn, table, schema=None):
     def get_columns(self, conn, table_name, schema=None, **kw):
         res = []
         # TODO(bdarnell): escape table name
-        for row in conn.execute('SHOW COLUMNS FROM "%s"' % table_name):
+        for row in conn.execute('SHOW COLUMNS FROM "%s"."%s"' % (schema or self.default_schema_name, table_name)):
             name, type_str, nullable, default = row[:4]
             # When there are type parameters, attach them to the
             # returned type object.
@@ -130,7 +132,7 @@ def get_indexes(self, conn, table_name, schema=None, **kw):
         uniques = collections.OrderedDict()
         columns = collections.defaultdict(list)
         # TODO(bdarnell): escape table name
-        for row in conn.execute('SHOW INDEXES FROM "%s"' % table_name):
+        for row in conn.execute('SHOW INDEXES FROM "%s"."%s"' % (schema or self.default_schema_name, table_name)):
             # beta-20170112 and older versions do not have the Implicit column.
             if getattr(row, "Implicit", False):
                 continue

cockroachdb/sqlalchemy/stmt_compiler.py

@@ -0,0 +1,13 @@
+from sqlalchemy.dialects.postgresql.psycopg2 import PGCompiler_psycopg2, expression
+
+
+class CockroachCompiler(PGCompiler_psycopg2):
+
+    # TODO(bdarnell): remove this when cockroachdb/cockroach#17008 is fixed.
+    def returning_clause(self, stmt, returning_cols):
+        columns = [
+            '.'.join(self._label_select_column(None, c, True, False, {}).split('.')[-2:])
+            for c in expression._select_iterables(returning_cols)
+        ]
+
+        return 'RETURNING ' + ', '.join(columns)

test/sqlalchemy/test_across_schema.py

@@ -0,0 +1,26 @@
+from sqlalchemy import Table, MetaData, testing
+from sqlalchemy.testing import fixtures
+
+
+class AcrossSchemaTest(fixtures.TestBase):
+    TEST_DATABASE = 'test_sqlalchemy_across_schema'
+
+    def teardown_method(self, method):
+        testing.db.execute("DROP TABLE IF EXISTS {}.users".format(self.TEST_DATABASE))
+        testing.db.execute("DROP DATABASE IF EXISTS {}".format(self.TEST_DATABASE))
+
+    def setup_method(self):
+        testing.db.execute("CREATE DATABASE IF NOT EXISTS {}".format(self.TEST_DATABASE))
+        testing.db.execute("CREATE TABLE IF NOT EXISTS {}.users (name STRING PRIMARY KEY)".format(self.TEST_DATABASE))
+        self.meta = MetaData(testing.db, schema=self.TEST_DATABASE)
+
+    def test_get_columns_indexes_across_schema(self):
+        # get_columns and get_indexes use default db uri schema.
+        # across schema table must use schema.table
+        Table('users', self.meta, autoload=True, schema=self.TEST_DATABASE)
+
+    def test_returning_clause(self):
+        # TODO(bdarnell): remove this when cockroachdb/cockroach#17008 is fixed.
+        # across schema returning is schema.table.id but cockroachdb not support.
+        table = Table('users', self.meta, autoload=True, schema=self.TEST_DATABASE)
+        table.insert().values(dict(name='John')).returning(table.c.name).execute()