Merge pull request #5654 from roed314/tablespaces

Add option for different tablespace to create_table

lmfdb/api/api.py

@@ -96,6 +96,7 @@ def split_db(tablename):
         dname, _ = split_db(tablename)
         dbSize[dname] += sizes['total_bytes']
         dbObjects[dname] += sizes['nrows']
+    tablespaces = db.tablespaces()
     for tablename, sizes in table_sizes.items():
         tsize = sizes['total_bytes']
         size += tsize
@@ -119,7 +120,9 @@ def split_db(tablename):
                 'size': csize, 'avgObjSize':avg_size,
                 'indexSize':mb(sizes['index_bytes']), 'dataSize':mb(sizes['table_bytes'] + sizes['toast_bytes'] + sizes['extras_bytes']),
                 'countsSize':mb(sizes['counts_bytes']), 'statsSize':mb(sizes['stats_bytes']),
-                'nrows': sizes['nrows'], 'nstats': sizes['nstats'], 'ncounts': sizes['ncounts']}
+                'nrows': sizes['nrows'], 'nstats': sizes['nstats'], 'ncounts': sizes['ncounts'],
+                'tablespace': tablespaces.get(tablename, ""),
+            }
     dataSize = size - indexSize
     info['ntables'] = len(table_sizes)
     info['nobjects'] = nobjects

lmfdb/api/templates/api-stats.html

@@ -52,6 +52,7 @@
   <th>Index (MiB)</th>
   <th>Objects</th>
   <th>Average Size (B)</th>
+  <th>Tablespace</th>
 </tr>
 </thead>
 
@@ -64,6 +65,7 @@
     <td>{{x.indexSize}}</td>
     <td>{{x.nrows}}</td>
     <td>{{x.avgObjSize}}</td>
+    <td>{{x.tablespace}}</td>
 </tr>
 {% endfor %}
 </tbody>

lmfdb/backend/base.py

@@ -908,6 +908,21 @@ def _copy_from(self, filename, table, columns, header, kwds):
 
                 return addid, cur.rowcount
 
+    def _get_tablespace(self):
+        # overridden in table and statstable
+        pass
+
+    def _tablespace_clause(self, tablespace=None):
+        """
+        A clause for use in CREATE statements
+        """
+        if tablespace is None:
+            tablespace = self._get_tablespace()
+        if tablespace is None:
+            return SQL("")
+        else:
+            return SQL(" TABLESPACE {0}").format(Identifier(tablespace))
+
     def _clone(self, table, tmp_table):
         """
         Utility function: creates a table with the same schema as the given one.
@@ -927,7 +942,7 @@ def _clone(self, table, tmp_table):
                 "Run db.%s.cleanup_from_reload() if you want to delete it and proceed."
                 % (tmp_table, table)
             )
-        creator = SQL("CREATE TABLE {0} (LIKE {1})").format(Identifier(tmp_table), Identifier(table))
+        creator = SQL("CREATE TABLE {0} (LIKE {1}){2}").format(Identifier(tmp_table), Identifier(table), self._tablespace_clause())
         self._execute(creator)
 
     def _check_col_datatype(self, typ):
@@ -937,7 +952,9 @@ def _check_col_datatype(self, typ):
 
     def _create_table(self, name, columns):
         """
-        Utility function: creates a table with the schema specified by ``columns``
+        Utility function: creates a table with the schema specified by ``columns``.
+
+        If self is a table, the new table will be in the same tablespace.
 
         INPUT:
 
@@ -949,7 +966,7 @@ def _create_table(self, name, columns):
         for col, typ in columns:
             self._check_col_datatype(typ)
         table_col = SQL(", ").join(SQL("{0} %s" % typ).format(Identifier(col)) for col, typ in columns)
-        creator = SQL("CREATE TABLE {0} ({1})").format(Identifier(name), table_col)
+        creator = SQL("CREATE TABLE {0} ({1}){2}").format(Identifier(name), table_col, self._tablespace_clause())
         self._execute(creator)
 
     def _create_table_from_header(self, filename, name, sep, addid=True):

lmfdb/backend/database.py

@@ -513,14 +513,17 @@ def _create_meta_tables_hist(self):
 
         print("Table meta_tables_hist created")
 
-    def create_table_like(self, new_name, table, data=False, commit=True):
+    def create_table_like(self, new_name, table, tablespace=None, data=False, indexes=False, commit=True):
         """
         Copies the schema from an existing table, but none of the data, indexes or stats.
 
         INPUT:
 
         - ``new_name`` -- a string giving the desired table name.
         - ``table`` -- a string or PostgresSearchTable object giving an existing table.
+        - ``tablespace`` -- the tablespace for the new table
+        - ``data`` -- whether to copy over data from the source table
+        - ``indexes`` -- whether to copy over indexes from the source table
         """
         if isinstance(table, str):
             table = self[table]
@@ -558,6 +561,7 @@ def create_table_like(self, new_name, table, data=False, commit=True):
             extra_columns,
             search_order,
             extra_order,
+            tablespace=tablespace,
             commit=commit,
         )
         if data:
@@ -577,6 +581,10 @@ def create_table_like(self, new_name, table, data=False, commit=True):
                     ),
                     commit=commit,
                 )
+        if indexes:
+            for idata in table.list_indexes(verbose=False).values():
+                self[new_name].create_index(**idata)
+        if data:
             self[new_name].stats.refresh_stats()
 
     def create_table(
@@ -591,6 +599,7 @@ def create_table(
         extra_columns=None,
         search_order=None,
         extra_order=None,
+        tablespace=None,
         force_description=False,
         commit=True,
     ):
@@ -618,6 +627,7 @@ def create_table(
             in the search table, speeding up scans.
         - ``search_order`` -- (optional) list of column names, specifying the default order of columns
         - ``extra_order`` -- (optional) list of column names, specifying the default order of columns
+        - ``tablespace`` -- (optional) a postgres tablespace to use for the new table
         - ``force_description`` -- whether to require descriptions
 
         COMMON TYPES:
@@ -729,35 +739,39 @@ def process_columns(coldict, colorder):
             if col_description is None:
                 col_description = {col: "" for col in description_columns}
 
+        tablespace = self._tablespace_clause(tablespace)
         with DelayCommit(self, commit, silence=True):
-            creator = SQL("CREATE TABLE {0} ({1})").format(
-                Identifier(name), SQL(", ").join(processed_search_columns)
+            creator = SQL("CREATE TABLE {0} ({1}){2}").format(
+                Identifier(name),
+                SQL(", ").join(processed_search_columns),
+                tablespace,
             )
             self._execute(creator)
             self.grant_select(name)
             if extra_columns is not None:
-                creator = SQL("CREATE TABLE {0} ({1})")
+                creator = SQL("CREATE TABLE {0} ({1}){2}")
                 creator = creator.format(
                     Identifier(name + "_extras"),
                     SQL(", ").join(processed_extra_columns),
+                    tablespace,
                 )
                 self._execute(creator)
                 self.grant_select(name + "_extras")
             creator = SQL(
                 "CREATE TABLE {0} "
                 "(cols jsonb, values jsonb, count bigint, "
-                "extra boolean, split boolean DEFAULT FALSE)"
+                "extra boolean, split boolean DEFAULT FALSE){1}"
             )
-            creator = creator.format(Identifier(name + "_counts"))
+            creator = creator.format(Identifier(name + "_counts"), tablespace)
             self._execute(creator)
             self.grant_select(name + "_counts")
             self.grant_insert(name + "_counts")
             creator = SQL(
                 "CREATE TABLE {0} "
                 '(cols jsonb, stat text COLLATE "C", value numeric, '
-                "constraint_cols jsonb, constraint_values jsonb, threshold integer)"
+                "constraint_cols jsonb, constraint_values jsonb, threshold integer){1}"
             )
-            creator = creator.format(Identifier(name + "_stats"))
+            creator = creator.format(Identifier(name + "_stats"), tablespace)
             self._execute(creator)
             self.grant_select(name + "_stats")
             self.grant_insert(name + "_stats")
@@ -1270,3 +1284,10 @@ def show_locks(self):
                 )
         else:
             print("No locks currently held")
+
+    def tablespaces(self):
+        """
+        Returns a dictionary giving giving the tablespace for all tables
+        """
+        D = {rec[0]: rec[1] for rec in self._execute(SQL("SELECT tablename, tablespace FROM pg_tables"))}
+        return {name: space if space else "" for (name, space) in D.items()}

lmfdb/backend/statstable.py

@@ -216,6 +216,10 @@ def __init__(self, table, total=None):
                 total = self._slow_count({}, extra=False)
         self.total = total
 
+    def _get_tablespace(self):
+        # We use the same tablespace for stats and counts tables as for the main search table
+        return self.table._get_tablespace()
+
     def _has_stats(self, jcols, ccols, cvals, threshold, split_list=False, threshold_inequality=False, suffix=""):
         """
         Checks whether statistics have been recorded for a given set of columns.

lmfdb/backend/table.py

@@ -30,7 +30,7 @@
         "varchar_ops",
         "varchar_pattern_ops",
     ],
-    "gin": ["jsonb_path_ops"],
+    "gin": ["jsonb_path_ops", "array_ops"],
     "gist": ["inet_ops"],
     "hash": [
         "bpchar_pattern_ops",
@@ -276,8 +276,14 @@ def list_indexes(self, verbose=False):
         if not verbose:
             return output
 
-    @staticmethod
-    def _create_index_statement(name, table, type, columns, modifiers, storage_params):
+    def _get_tablespace(self):
+        """
+        Determine the tablespace hosting this table (which is then used for indexes and constraints)
+        """
+        cur = self._execute(SQL("SELECT tablespace FROM pg_tables WHERE tablename=%s"), [self.search_table])
+        return cur.fetchone()[0]
+
+    def _create_index_statement(self, name, table, type, columns, modifiers, storage_params):
         """
         Utility function for making the create index SQL statement.
         """
@@ -291,15 +297,16 @@ def _create_index_statement(name, table, type, columns, modifiers, storage_param
             )
         else:
             storage_params = SQL("")
+        tablespace = self._tablespace_clause()
         modifiers = [" " + " ".join(mods) if mods else "" for mods in modifiers]
         # The inner % operator is on strings prior to being wrapped by SQL: modifiers have been whitelisted.
         columns = SQL(", ").join(
             SQL("{0}%s" % mods).format(Identifier(col))
             for col, mods in zip(columns, modifiers)
         )
         # The inner % operator is on strings prior to being wrapped by SQL: type has been whitelisted.
-        creator = SQL("CREATE INDEX {0} ON {1} USING %s ({2}){3}" % (type))
-        return creator.format(Identifier(name), Identifier(table), columns, storage_params)
+        creator = SQL("CREATE INDEX {0} ON {1} USING %s ({2}){3}{4}" % (type))
+        return creator.format(Identifier(name), Identifier(table), columns, storage_params, tablespace)
 
     def _create_counts_indexes(self, suffix="", warning_only=False):
         """
@@ -447,6 +454,13 @@ def mod(col):
                 name = "_".join([self.search_table] + [col[:2] for col in columns])
             else:
                 name = "_".join([self.search_table] + ["".join(col[0] for col in columns)])
+            if len(name) >= 64:
+                name = name[:63]
+            if self._relation_exists(name):
+                disamb = 0
+                while self._relation_exists(name + str(disamb)):
+                    disamb += 1
+                name += str(disamb)
 
         with DelayCommit(self, silence=True):
             self._check_index_name(name, "Index")
@@ -1080,7 +1094,7 @@ def drop_tmp():
                 SQL("{0} " + self.col_type[col]).format(Identifier(col))
                 for col in columns
             ])
-            creator = SQL("CREATE TABLE {0} ({1})").format(Identifier(tmp_table), processed_columns)
+            creator = SQL("CREATE TABLE {0} ({1}){2}").format(Identifier(tmp_table), processed_columns, self._tablespace_clause())
             self._execute(creator)
             # We need to add an id column and populate it correctly
             if label_col != "id":
@@ -2435,7 +2449,7 @@ def create_extra_table(self, columns, ordered=False, sep="|", commit=True):
             col_type_SQL = SQL(", ").join(
                 SQL("{0} %s" % typ).format(Identifier(col)) for col, typ in col_type
             )
-            creator = SQL("CREATE TABLE {0} ({1})").format(Identifier(self.extra_table), col_type_SQL)
+            creator = SQL("CREATE TABLE {0} ({1}){2}").format(Identifier(self.extra_table), col_type_SQL, self._tablespace_clause())
             self._execute(creator)
             if columns:
                 self.drop_constraints(columns)