database_observability: add foreign keys to PG schema_details collector (#4549)

fridgepoet · web-flow · commit e02c9008947a · 2025-10-07T05:52:54.000-07:00
diff --git a/internal/component/database_observability/postgres/collector/schema_details.go b/internal/component/database_observability/postgres/collector/schema_details.go
@@ -126,6 +126,38 @@ const (
 	GROUP BY index_relations.relname, pg_am.amname, pg_index.indisunique
 	ORDER BY index_name
 `
+
+	// selectForeignKeys retrieves foreign key constraints for a specified table
+	/*
+		pg_constraint stores all constraints
+		join pg_class (table info) to get the source table
+		join to pg_namespace (schema info) for schema filtering
+		join to pg_class again to get referenced table
+		use generate_subscripts() to correlate multi-column foreign keys by position
+		pg_attribute joined twice to get column names for both source and referenced columns
+	*/
+	selectForeignKeys = `
+	SELECT
+		constraints.conname as constraint_name,
+		source_column.attname as column_name,
+		referenced_table.relname as referenced_table_name,
+		referenced_column.attname as referenced_column_name
+	FROM pg_constraint constraints
+	JOIN pg_class source_table ON constraints.conrelid = source_table.oid
+	JOIN pg_namespace schema ON source_table.relnamespace = schema.oid
+	JOIN pg_class referenced_table ON constraints.confrelid = referenced_table.oid
+	JOIN generate_subscripts(constraints.conkey, 1) AS position ON true
+	JOIN pg_attribute source_column ON constraints.conrelid = source_column.attrelid
+		AND source_column.attnum = constraints.conkey[position]
+		AND NOT source_column.attisdropped
+	JOIN pg_attribute referenced_column ON constraints.confrelid = referenced_column.attrelid
+		AND referenced_column.attnum = constraints.confkey[position]
+		AND NOT referenced_column.attisdropped
+	WHERE constraints.contype = 'f'
+		AND schema.nspname = $1
+		AND source_table.relname = $2
+	ORDER BY constraints.conname, position
+`
 )
 
 type tableInfo struct {
@@ -137,8 +169,9 @@ type tableInfo struct {
 }
 
 type tableSpec struct {
-	Columns []columnSpec `json:"columns"`
-	Indexes []indexSpec  `json:"indexes,omitempty"`
+	Columns     []columnSpec `json:"columns"`
+	Indexes     []indexSpec  `json:"indexes,omitempty"`
+	ForeignKeys []foreignKey `json:"foreign_keys,omitempty"`
 }
 
 type columnSpec struct {
@@ -159,6 +192,13 @@ type indexSpec struct {
 	Nullable    bool     `json:"nullable"`
 }
 
+type foreignKey struct {
+	Name                 string `json:"name"`
+	ColumnName           string `json:"column_name"`
+	ReferencedTableName  string `json:"referenced_table_name"`
+	ReferencedColumnName string `json:"referenced_column_name"`
+}
+
 type SchemaDetailsArguments struct {
 	DB           *sql.DB
 	EntryHandler loki.EntryHandler
@@ -430,5 +470,32 @@ func (c *SchemaDetails) fetchColumnsDefinitions(ctx context.Context, databaseNam
 		return nil, err
 	}
 
+	fkRS, err := c.dbConnection.QueryContext(ctx, selectForeignKeys, schemaName, tableName)
+	if err != nil {
+		level.Error(c.logger).Log("msg", "failed to query foreign keys", "datname", databaseName, "schema", schemaName, "table", tableName, "err", err)
+		return nil, err
+	}
+	defer fkRS.Close()
+
+	for fkRS.Next() {
+		var constraintName, columnName, referencedTableName, referencedColumnName string
+		if err := fkRS.Scan(&constraintName, &columnName, &referencedTableName, &referencedColumnName); err != nil {
+			level.Error(c.logger).Log("msg", "failed to scan foreign keys", "datname", databaseName, "schema", schemaName, "table", tableName, "err", err)
+			return nil, err
+		}
+
+		tblSpec.ForeignKeys = append(tblSpec.ForeignKeys, foreignKey{
+			Name:                 constraintName,
+			ColumnName:           columnName,
+			ReferencedTableName:  referencedTableName,
+			ReferencedColumnName: referencedColumnName,
+		})
+	}
+
+	if err := fkRS.Err(); err != nil {
+		level.Error(c.logger).Log("msg", "failed to iterate over foreign keys result set", "datname", databaseName, "schema", schemaName, "table", tableName, "err", err)
+		return nil, err
+	}
+
 	return tblSpec, nil
 }
diff --git a/internal/component/database_observability/postgres/collector/schema_details_test.go b/internal/component/database_observability/postgres/collector/schema_details_test.go
@@ -90,6 +90,16 @@ func TestSchemaTable(t *testing.T) {
 				}).AddRow("authors_pkey", "btree", true, pq.StringArray{"id"}, pq.StringArray{}, true),
 			)
 
+		mock.ExpectQuery(selectForeignKeys).WithArgs("public", "authors").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"constraint_name",
+					"column_name",
+					"referenced_table_name",
+					"referenced_column_name",
+				}),
+			)
+
 		err = collector.Start(t.Context())
 		require.NoError(t, err)
 
@@ -189,6 +199,16 @@ func TestSchemaTable(t *testing.T) {
 				}).AddRow("authors_pkey", "btree", true, pq.StringArray{"id"}, pq.StringArray{}, false),
 			)
 
+		mock.ExpectQuery(selectForeignKeys).WithArgs("public", "authors").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"constraint_name",
+					"column_name",
+					"referenced_table_name",
+					"referenced_column_name",
+				}),
+			)
+
 		mock.ExpectQuery(selectColumnNames).WithArgs("public.categories").RowsWillBeClosed().
 			WillReturnRows(
 				sqlmock.NewRows([]string{
@@ -213,6 +233,16 @@ func TestSchemaTable(t *testing.T) {
 				}).AddRow("categories_pkey", "btree", true, pq.StringArray{"id"}, pq.StringArray{}, false),
 			)
 
+		mock.ExpectQuery(selectForeignKeys).WithArgs("public", "categories").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"constraint_name",
+					"column_name",
+					"referenced_table_name",
+					"referenced_column_name",
+				}),
+			)
+
 		mock.ExpectQuery(selectColumnNames).WithArgs("postgis.spatial_ref_sys").RowsWillBeClosed().
 			WillReturnRows(
 				sqlmock.NewRows([]string{
@@ -237,6 +267,16 @@ func TestSchemaTable(t *testing.T) {
 				}),
 			)
 
+		mock.ExpectQuery(selectForeignKeys).WithArgs("postgis", "spatial_ref_sys").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"constraint_name",
+					"column_name",
+					"referenced_table_name",
+					"referenced_column_name",
+				}),
+			)
+
 		err = collector.Start(t.Context())
 		require.NoError(t, err)
 
@@ -346,6 +386,16 @@ func TestSchemaTable(t *testing.T) {
 					AddRow("idx_users_created_at", "btree", false, pq.StringArray{"created_at"}, nil, false),
 			)
 
+		mock.ExpectQuery(selectForeignKeys).WithArgs("public", "users").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"constraint_name",
+					"column_name",
+					"referenced_table_name",
+					"referenced_column_name",
+				}),
+			)
+
 		err = collector.Start(t.Context())
 		require.NoError(t, err)
 
@@ -356,6 +406,7 @@ func TestSchemaTable(t *testing.T) {
 		collector.Stop()
 		lokiClient.Stop()
 
+		// Run this after Stop() to avoid race conditions
 		err = mock.ExpectationsWereMet()
 		require.NoError(t, err)
 
@@ -487,6 +538,16 @@ func TestSchemaTable(t *testing.T) {
 				}),
 			)
 
+		mock.ExpectQuery(selectForeignKeys).WithArgs("public", "test_table").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"constraint_name",
+					"column_name",
+					"referenced_table_name",
+					"referenced_column_name",
+				}),
+			)
+
 		err = collector.Start(t.Context())
 		require.NoError(t, err)
 
@@ -581,6 +642,16 @@ func Test_collector_detects_auto_increment_column(t *testing.T) {
 				}),
 			)
 
+		mock.ExpectQuery(selectForeignKeys).WithArgs("public", "users").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"constraint_name",
+					"column_name",
+					"referenced_table_name",
+					"referenced_column_name",
+				}),
+			)
+
 		err = collector.Start(t.Context())
 		require.NoError(t, err)
 
@@ -672,6 +743,16 @@ func Test_collector_detects_auto_increment_column(t *testing.T) {
 				}),
 			)
 
+		mock.ExpectQuery(selectForeignKeys).WithArgs("public", "products").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"constraint_name",
+					"column_name",
+					"referenced_table_name",
+					"referenced_column_name",
+				}),
+			)
+
 		err = collector.Start(t.Context())
 		require.NoError(t, err)
 
@@ -696,4 +777,101 @@ func Test_collector_detects_auto_increment_column(t *testing.T) {
 		expectedTableSpec := base64.StdEncoding.EncodeToString([]byte(`{"columns":[{"name":"id","type":"integer","not_null":true,"auto_increment":true,"primary_key":true},{"name":"code","type":"integer","not_null":true,"auto_increment":true},{"name":"name","type":"character varying(255)","not_null":true}]}`))
 		require.Equal(t, fmt.Sprintf(`level="info" datname="identity_test_db" schema="public" table="products" table_spec="%s"`, expectedTableSpec), lokiEntries[2].Line)
 	})
+
+	t.Run("collector detects foreign keys", func(t *testing.T) {
+		t.Parallel()
+
+		db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
+		require.NoError(t, err)
+		defer db.Close()
+
+		lokiClient := loki_fake.NewClient(func() {})
+
+		collector, err := NewSchemaDetails(SchemaDetailsArguments{
+			DB:           db,
+			EntryHandler: lokiClient,
+			Logger:       log.NewLogfmtLogger(os.Stderr),
+		})
+		require.NoError(t, err)
+		require.NotNil(t, collector)
+
+		mock.ExpectQuery(selectDatabaseName).WithoutArgs().RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"datname",
+				}).AddRow(
+					"books_store",
+				),
+			)
+
+		mock.ExpectQuery(selectSchemaNames).WithoutArgs().RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"schema_name",
+				}).AddRow("public"),
+			)
+
+		mock.ExpectQuery(selectTableNames).WithArgs("public").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"table_name",
+				}).AddRow("books"),
+			)
+
+		mock.ExpectQuery(selectColumnNames).WithArgs("public.books").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"column_name",
+					"column_type",
+					"not_nullable",
+					"column_default",
+					"identity_generation",
+					"is_primary_key",
+				}).AddRow("id", "integer", true, "", "", true).
+					AddRow("title", "character varying(255)", true, "", "", false).
+					AddRow("author_id", "integer", true, "", "", false).
+					AddRow("category_id", "integer", false, "", "", false),
+			)
+
+		mock.ExpectQuery(selectIndexes).WithArgs("public", "books").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"index_name",
+					"index_type",
+					"unique",
+					"column_names",
+					"expressions",
+					"has_nullable_column",
+				}).AddRow("books_pkey", "btree", true, pq.StringArray{"id"}, pq.StringArray{}, false),
+			)
+
+		mock.ExpectQuery(selectForeignKeys).WithArgs("public", "books").RowsWillBeClosed().
+			WillReturnRows(
+				sqlmock.NewRows([]string{
+					"constraint_name",
+					"column_name",
+					"referenced_table_name",
+					"referenced_column_name",
+				}).AddRow("fk_books_author", "author_id", "authors", "id").
+					AddRow("fk_books_category", "category_id", "categories", "id"),
+			)
+
+		err = collector.Start(t.Context())
+		require.NoError(t, err)
+
+		require.Eventually(t, func() bool {
+			return len(lokiClient.Received()) == 3
+		}, 2*time.Second, 100*time.Millisecond)
+
+		collector.Stop()
+		lokiClient.Stop()
+
+		err = mock.ExpectationsWereMet()
+		require.NoError(t, err)
+
+		lokiEntries := lokiClient.Received()
+		require.Len(t, lokiEntries, 3)
+		expectedTableSpec := base64.StdEncoding.EncodeToString([]byte(`{"columns":[{"name":"id","type":"integer","not_null":true,"primary_key":true},{"name":"title","type":"character varying(255)","not_null":true},{"name":"author_id","type":"integer","not_null":true},{"name":"category_id","type":"integer"}],"indexes":[{"name":"books_pkey","type":"btree","columns":["id"],"unique":true,"nullable":false}],"foreign_keys":[{"name":"fk_books_author","column_name":"author_id","referenced_table_name":"authors","referenced_column_name":"id"},{"name":"fk_books_category","column_name":"category_id","referenced_table_name":"categories","referenced_column_name":"id"}]}`))
+		require.Equal(t, fmt.Sprintf(`level="info" datname="books_store" schema="public" table="books" table_spec="%s"`, expectedTableSpec), lokiEntries[2].Line)
+	})
 }
